Load the RoRF router and embedding model from the specified local or Huggingface repo.
(self, repo_id: str, embedding_provider: str)
| 104 | return recommended_model, model_a_prob, model_b_prob |
| 105 | |
| 106 | def load(self, repo_id: str, embedding_provider: str): |
| 107 | """ |
| 108 | Load the RoRF router and embedding model from the specified local or Huggingface repo. |
| 109 | """ |
| 110 | logger.info(f"Initializing {embedding_provider} embedding model...") |
| 111 | embedding_model = get_embedding_model(embedding_provider) |
| 112 | |
| 113 | if os.path.exists(f"{repo_id}/classifier.pkl"): |
| 114 | logger.info(f"{repo_id} exists. Loading router...") |
| 115 | router_path = f"{repo_id}/classifier.pkl" |
| 116 | else: |
| 117 | logger.info(f"{repo_id} does not exist. Downloading router...") |
| 118 | router_path = hf_hub_download( |
| 119 | repo_id=repo_id, |
| 120 | filename="classifier.pkl", |
| 121 | ) |
| 122 | with open(router_path, "rb") as f: |
| 123 | router_model = pickle.load(f) |
| 124 | return router_model, embedding_model |
no test coverage detected