Given a prompt or prompt embedding, return the model recommended by the RoRF router along with the probabilities of the two models.
(self, prompt: Union[str, np.ndarray])
| 86 | return recommended_model |
| 87 | |
| 88 | def predict_proba(self, prompt: Union[str, np.ndarray]) -> str: |
| 89 | """ |
| 90 | Given a prompt or prompt embedding, return the model recommended by the RoRF router along with the probabilities of the two models. |
| 91 | """ |
| 92 | if isinstance(prompt, str): |
| 93 | prompt_embedding = self.embedding_model.get_prompt_embeddings([prompt]) |
| 94 | else: |
| 95 | prompt_embedding = prompt |
| 96 | label_scores = self.router_model.predict_proba(prompt_embedding)[0] |
| 97 | model_a_prob = label_scores[0] + label_scores[1] |
| 98 | model_b_prob = label_scores[2] + label_scores[3] |
| 99 | |
| 100 | if model_a_prob >= self.threshold: |
| 101 | recommended_model = self.model_a |
| 102 | else: |
| 103 | recommended_model = self.model_b |
| 104 | return recommended_model, model_a_prob, model_b_prob |
| 105 | |
| 106 | def load(self, repo_id: str, embedding_provider: str): |
| 107 | """ |
no test coverage detected