Spaces:
Running
Running
| from __future__ import annotations | |
| from services import hf_provider, featherless_provider | |
| class ModelRouter: | |
| """ | |
| Routes tasks to the correct model and provider. | |
| understand() → MiniCPM-V (HF) — all document tasks | |
| reason() → Nemotron 4B (Featherless) — all reasoning tasks | |
| translate() → Tiny Aya (HF) — multilingual | |
| transcribe() → Whisper (HF) — speech to text | |
| """ | |
| def __init__(self, ocr_model: str, reasoning_model: str, | |
| multilingual_model: str, speech_model: str, | |
| hf_api_key: str = "", featherless_api_key: str = "", | |
| max_tokens: int = 1024, temperature: float = 0.3): | |
| self._ocr_model = ocr_model | |
| self._reason_model = reasoning_model | |
| self._multi_model = multilingual_model | |
| self._speech_model = speech_model | |
| self._hf_key = hf_api_key | |
| self._fl_key = featherless_api_key | |
| self._max_tokens = max_tokens | |
| self._temperature = temperature | |
| def understand(self, prompt: str, image_b64: str = "") -> str: | |
| """MiniCPM-V via HuggingFace — OCR, diagram reading, concept extraction.""" | |
| if image_b64: | |
| return hf_provider.vision_generate( | |
| self._ocr_model, image_b64, prompt, | |
| self._hf_key, self._max_tokens) | |
| return hf_provider.generate( | |
| self._ocr_model, prompt, api_key=self._hf_key, | |
| max_tokens=self._max_tokens, temperature=self._temperature) | |
| def reason(self, prompt: str, system: str = "") -> str: | |
| """Nemotron 3 Nano 4B via Featherless — quests, questions, tutor.""" | |
| return featherless_provider.generate( | |
| self._reason_model, prompt, system, | |
| self._fl_key, self._max_tokens, self._temperature) | |
| def translate(self, prompt: str, system: str = "") -> str: | |
| """Tiny Aya 3.3B via HuggingFace — multilingual explanations.""" | |
| return hf_provider.generate( | |
| self._multi_model, prompt, system, | |
| self._hf_key, self._max_tokens, self._temperature) | |
| def transcribe(self, audio_bytes: bytes) -> str: | |
| """Whisper via HuggingFace — speech to text.""" | |
| return hf_provider.transcribe(self._speech_model, audio_bytes, self._hf_key) | |