Load the LLM router model. returns: AdaptiveClassifier: The loaded model exceptions: Exception: If the safetensors fails to load
(self)
| 43 | } |
| 44 | |
| 45 | def load_llm_router(self) -> AdaptiveClassifier: |
| 46 | """ |
| 47 | Load the LLM router model. |
| 48 | returns: |
| 49 | AdaptiveClassifier: The loaded model |
| 50 | exceptions: |
| 51 | Exception: If the safetensors fails to load |
| 52 | """ |
| 53 | path = "../llm_router" if __name__ == "__main__" else "./llm_router" |
| 54 | try: |
| 55 | animate_thinking("Loading LLM router model...", color="status") |
| 56 | talk_classifier = AdaptiveClassifier.from_pretrained(path) |
| 57 | except Exception as e: |
| 58 | raise Exception("Failed to load the routing model. Please run the dl_safetensors.sh script inside llm_router/ directory to download the model.") |
| 59 | return talk_classifier |
| 60 | |
| 61 | def get_device(self) -> str: |
| 62 | if torch.backends.mps.is_available(): |
no test coverage detected