Browse Source

Add weighted random strategy

pull/2261/head
Hironsan 1 year ago
parent
commit
b4c72af5a3
1 changed files with 29 additions and 0 deletions
  1. 29
      backend/examples/assignment/strategies.py

29
backend/examples/assignment/strategies.py

@ -0,0 +1,29 @@
import abc
import dataclasses
from typing import List
import numpy as np
@dataclasses.dataclass
class Assignment:
user: int
example: int
class BaseStrategy(abc.ABC):
@abc.abstractmethod
def assign(self) -> List[Assignment]:
...
class WeightedRandomStrategy:
def __init__(self, dataset_size: int, weights: List[int]):
assert sum(weights) == 100
self.dataset_size = dataset_size
self.weights = weights
def assign(self) -> List[Assignment]:
proba = np.array(self.weights) / 100
assignees = np.random.choice(range(len(self.weights)), size=self.dataset_size, p=proba)
return [Assignment(user=user, example=example) for example, user in enumerate(assignees)]
Loading…
Cancel
Save