| 21 | |
| 22 | |
| 23 | class Controller: |
| 24 | def __init__( |
| 25 | self, |
| 26 | router: str, |
| 27 | model_a: str, |
| 28 | model_b: str, |
| 29 | threshold: float, |
| 30 | ): |
| 31 | """ |
| 32 | Initialize the RoRF controller with the specified router, model A, model B, and threshold. |
| 33 | Threshold determines the percentage of calls made to model A by the router. |
| 34 | """ |
| 35 | self._validate_router_threshold(router, threshold) |
| 36 | self.embedding_provider = self._parse_model_name(router) |
| 37 | logger.info(f"Initializing RoRF controller for {router} with {self.embedding_provider} embeddings...") |
| 38 | self.router_model, self.embedding_model = self.load(router, self.embedding_provider) |
| 39 | self.model_a, self.model_b = model_a, model_b |
| 40 | self.threshold = threshold |
| 41 | |
| 42 | def _validate_router_threshold( |
| 43 | self, router: Optional[str], threshold: Optional[float] |
| 44 | ): |
| 45 | """ |
| 46 | Validate the router and threshold. |
| 47 | """ |
| 48 | if router is None or threshold is None: |
| 49 | raise RoutingError("Router or threshold unspecified.") |
| 50 | if not 0 <= threshold <= 1: |
| 51 | raise RoutingError( |
| 52 | f"Invalid threshold {threshold}. Threshold must be a float between 0.0 and 1.0." |
| 53 | ) |
| 54 | |
| 55 | def _parse_model_name(self, router: str): |
| 56 | """ |
| 57 | Parse the method and embedding model provider's name from the router name. |
| 58 | """ |
| 59 | method, embedding_provider = router.split("/")[1].split('-')[0:2] |
| 60 | if not method == "rorf": |
| 61 | raise RoutingError(f"Invalid method {method}. Method must be 'rorf'.") |
| 62 | if not embedding_provider in ["jina", "voyage", "openai"]: |
| 63 | raise RoutingError( |
| 64 | f"Invalid embedding provider {embedding_provider}. Embedding provider must be 'jina', 'voyage', or 'openai'." |
| 65 | ) |
| 66 | return embedding_provider |
| 67 | |
| 68 | def batch_calculate_win_rate(self, prompts: List[str]) -> List[float]: |
| 69 | """ |
| 70 | Given a list of prompts, calculate the win rates (Model A probability) using the RoRF router. |
| 71 | """ |
| 72 | logger.info(f"Calculating win rates for {len(prompts)}...") |
| 73 | model_a_probs = [] |
| 74 | prompt_embeddings = self.embedding_model.get_prompt_embeddings(prompts) |
| 75 | for prompt in tqdm(prompt_embeddings, total=len(prompt_embeddings)): |
| 76 | prompt = prompt.reshape((1, -1)) |
| 77 | _, model_a_prob, _ = self.predict_proba(prompt) |
| 78 | model_a_probs.append(model_a_prob) |
| 79 | return model_a_probs |
| 80 |
no outgoing calls
no test coverage detected