fix-changes

Dockerfile

@@ -6,6 +6,5 @@ COPY . .
 ENV HF_HOME=/data/huggingface
 # Expose port
 EXPOSE 7860
-RUN pip install torchvision
 # Start the server
 CMD ["python", "/app/src/server/main.py", "--host", "0.0.0.0", "--port", "7860", "--config", "config_two"]

Dockerfile.base

@@ -29,7 +29,6 @@ COPY requirements.txt .
 
 # Install Python dependencies
 RUN pip install --no-cache-dir -r requirements.txt
-
 # Set up user
 RUN useradd -ms /bin/bash appuser \
     && chown -R appuser:appuser /app

requirements.txt

@@ -176,7 +176,7 @@ torch==2.6.0
 torchaudio==2.6.0
 torchdiffeq==0.2.5
 tqdm==4.67.1
-transformers
+transformers==4.50.3
 transformers-stream-generator==0.0.5
 triton==3.2.0
 typer==0.15.2

src/server/main.py

@@ -387,68 +387,61 @@ SUPPORTED_LANGUAGES = {
 class TranslateManager:
     def __init__(self, src_lang, tgt_lang, device_type=device, use_distilled=True):
         self.device_type = device_type
-        self.tokenizer = None
-        self.model = None
-        self.src_lang = src_lang
-        self.tgt_lang = tgt_lang
-        self.use_distilled = use_distilled
+        self.tokenizer, self.model = self.initialize_model(src_lang, tgt_lang, use_distilled)
 
-    def load(self):
-        if not self.tokenizer or not self.model:
-            if self.src_lang.startswith("eng") and not self.tgt_lang.startswith("eng"):
-                model_name = "ai4bharat/indictrans2-en-indic-dist-200M" if self.use_distilled else "ai4bharat/indictrans2-en-indic-1B"
-            elif not self.src_lang.startswith("eng") and self.tgt_lang.startswith("eng"):
-                model_name = "ai4bharat/indictrans2-indic-en-dist-200M" if self.use_distilled else "ai4bharat/indictrans2-indic-en-1B"
-            elif not self.src_lang.startswith("eng") and not self.tgt_lang.startswith("eng"):
-                model_name = "ai4bharat/indictrans2-indic-indic-dist-320M" if self.use_distilled else "ai4bharat/indictrans2-indic-indic-1B"
-            else:
-                raise ValueError("Invalid language combination")
+    def initialize_model(self, src_lang, tgt_lang, use_distilled):
+        if src_lang.startswith("eng") and not tgt_lang.startswith("eng"):
+            model_name = "ai4bharat/indictrans2-en-indic-dist-200M" if use_distilled else "ai4bharat/indictrans2-en-indic-1B"
+        elif not src_lang.startswith("eng") and tgt_lang.startswith("eng"):
+            model_name = "ai4bharat/indictrans2-indic-en-dist-200M" if use_distilled else "ai4bharat/indictrans2-indic-en-1B"
+        elif not src_lang.startswith("eng") and not tgt_lang.startswith("eng"):
+            model_name = "ai4bharat/indictrans2-indic-indic-dist-320M" if use_distilled else "ai4bharat/indictrans2-indic-indic-1B"
+        else:
+            raise ValueError("Invalid language combination: English to English translation is not supported.")
 
-            self.tokenizer = AutoTokenizer.from_pretrained(
-                model_name,
-                trust_remote_code=True
-            )
-            self.model = AutoModelForSeq2SeqLM.from_pretrained(
-                model_name,
-                trust_remote_code=True,
-                torch_dtype=torch.float16,
-                attn_implementation="flash_attention_2"
-            )
-            self.model = self.model.to(self.device_type)
-            self.model = torch.compile(self.model, mode="reduce-overhead")
-            logger.info(f"Translation model {model_name} loaded")
+        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+        model = AutoModelForSeq2SeqLM.from_pretrained(
+            model_name,
+            trust_remote_code=True,
+            torch_dtype=torch.float16,
+            attn_implementation="flash_attention_2"
+        ).to(self.device_type)
+        return tokenizer, model
 
 class ModelManager:
     def __init__(self, device_type=device, use_distilled=True, is_lazy_loading=False):
-        self.models = {}
+        self.models: dict[str, TranslateManager] = {}
         self.device_type = device_type
         self.use_distilled = use_distilled
         self.is_lazy_loading = is_lazy_loading
+        # Preload all translation models
+        self.preload_models()
+
+    def preload_models(self):
+        # Define the core translation pairs to preload
+        translation_pairs = [
+            ('eng_Latn', 'kan_Knda', 'eng_indic'),  # English to Indic
+            ('kan_Knda', 'eng_Latn', 'indic_eng'),  # Indic to English
+            ('kan_Knda', 'hin_Deva', 'indic_indic')  # Indic to Indic
+        ]
+        for src_lang, tgt_lang, key in translation_pairs:
+            logger.info(f"Preloading translation model for {src_lang} -> {tgt_lang}...")
+            self.models[key] = TranslateManager(src_lang, tgt_lang, self.device_type, self.use_distilled)
+            logger.info(f"Translation model for {key} preloaded successfully")
 
-    def load_model(self, src_lang, tgt_lang, key):
-        logger.info(f"Loading translation model for {src_lang} -> {tgt_lang}")
-        translate_manager = TranslateManager(src_lang, tgt_lang, self.device_type, self.use_distilled)
-        translate_manager.load()
-        self.models[key] = translate_manager
-        logger.info(f"Loaded translation model for {key}")
-
-    def get_model(self, src_lang, tgt_lang):
-        key = self._get_model_key(src_lang, tgt_lang)
-        if key not in self.models:
-            if self.is_lazy_loading:
-                self.load_model(src_lang, tgt_lang, key)
-            else:
-                raise ValueError(f"Model for {key} is not preloaded and lazy loading is disabled.")
-        return self.models.get(key)
-
-    def _get_model_key(self, src_lang, tgt_lang):
+    def get_model(self, src_lang, tgt_lang) -> TranslateManager:
         if src_lang.startswith("eng") and not tgt_lang.startswith("eng"):
-            return 'eng_indic'
+            key = 'eng_indic'
         elif not src_lang.startswith("eng") and tgt_lang.startswith("eng"):
-            return 'indic_eng'
+            key = 'indic_eng'
         elif not src_lang.startswith("eng") and not tgt_lang.startswith("eng"):
-            return 'indic_indic'
-        raise ValueError("Invalid language combination")
+            key = 'indic_indic'
+        else:
+            raise ValueError("Invalid language combination: English to English translation is not supported.")
+
+        if key not in self.models:
+            raise ValueError(f"Model for {key} is not preloaded. All models should be preloaded at startup.")
+        return self.models[key]
 
 # ASR Manager
 class ASRModelManager:
@@ -510,6 +503,41 @@ class TranslationResponse(BaseModel):
 def get_translate_manager(src_lang: str, tgt_lang: str) -> TranslateManager:
     return model_manager.get_model(src_lang, tgt_lang)
 
+# Translation Function
+async def perform_internal_translation(sentences: List[str], src_lang: str, tgt_lang: str) -> List[str]:
+    translate_manager = model_manager.get_model(src_lang, tgt_lang)
+    if not sentences:
+        raise HTTPException(status_code=400, detail="Input sentences are required")
+
+    batch = ip.preprocess_batch(sentences, src_lang=src_lang, tgt_lang=tgt_lang)
+    inputs = translate_manager.tokenizer(
+        batch,
+        truncation=True,
+        padding="longest",
+        return_tensors="pt",
+        return_attention_mask=True,
+    ).to(translate_manager.device_type)
+
+    with torch.no_grad():
+        generated_tokens = translate_manager.model.generate(
+            **inputs,
+            use_cache=True,
+            min_length=0,
+            max_length=256,
+            num_beams=5,
+            num_return_sequences=1,
+        )
+
+    with translate_manager.tokenizer.as_target_tokenizer():
+        generated_tokens = translate_manager.tokenizer.batch_decode(
+            generated_tokens.detach().cpu().tolist(),
+            skip_special_tokens=True,
+            clean_up_tokenization_spaces=True,
+        )
+
+    translations = ip.postprocess_batch(generated_tokens, lang=tgt_lang)
+    return translations
+
 # Lifespan Event Handler
 translation_configs = []
 
@@ -532,23 +560,8 @@ async def lifespan(app: FastAPI):
             asr_manager.load()
             logger.info("ASR model loaded successfully")
 
-            # Load translation models
-            translation_tasks = [
-                ('eng_Latn', 'kan_Knda', 'eng_indic'),
-                ('kan_Knda', 'eng_Latn', 'indic_eng'),
-                ('kan_Knda', 'hin_Deva', 'indic_indic'),
-            ]
-            
-            for config in translation_configs:
-                src_lang = config["src_lang"]
-                tgt_lang = config["tgt_lang"]
-                key = model_manager._get_model_key(src_lang, tgt_lang)
-                translation_tasks.append((src_lang, tgt_lang, key))
-
-            for src_lang, tgt_lang, key in translation_tasks:
-                logger.info(f"Loading translation model for {src_lang} -> {tgt_lang}...")
-                model_manager.load_model(src_lang, tgt_lang, key)
-                logger.info(f"Translation model for {key} loaded successfully")
+            # Translation models are preloaded in ModelManager constructor
+            logger.info("Translation models already preloaded in ModelManager initialization.")
 
             logger.info("All models loaded successfully")
         except Exception as e:
@@ -625,7 +638,6 @@ async def chat(request: Request, chat_request: ChatRequest):
         raise HTTPException(status_code=400, detail="Prompt cannot be empty")
     logger.info(f"Received prompt: {chat_request.prompt}, src_lang: {chat_request.src_lang}, tgt_lang: {chat_request.tgt_lang}")
     try:
-        # Step 1: Translate prompt to English if needed
         if chat_request.src_lang != "eng_Latn":
             translated_prompt = await perform_internal_translation(
                 sentences=[chat_request.prompt],
@@ -638,11 +650,9 @@ async def chat(request: Request, chat_request: ChatRequest):
             prompt_to_process = chat_request.prompt
             logger.info("Prompt already in English, no translation needed")
 
-        # Step 2: Generate response in English
         response = await llm_manager.generate(prompt_to_process, settings.max_tokens)
         logger.info(f"Generated English response: {response}")
 
-        # Step 3: Translate response to target language if needed
         if chat_request.tgt_lang != "eng_Latn":
             translated_response = await perform_internal_translation(
                 sentences=[response],
@@ -672,7 +682,6 @@ async def visual_query(
         if image.size == (0, 0):
             raise HTTPException(status_code=400, detail="Uploaded image is empty or invalid")
 
-        # Step 1: Translate query to English if needed
         if src_lang != "eng_Latn":
             translated_query = await perform_internal_translation(
                 sentences=[query],
@@ -685,11 +694,9 @@ async def visual_query(
             query_to_process = query
             logger.info("Query already in English, no translation needed")
 
-        # Step 2: Generate answer in English
         answer = await llm_manager.vision_query(image, query_to_process)
         logger.info(f"Generated English answer: {answer}")
 
-        # Step 3: Translate answer to target language if needed
         if tgt_lang != "eng_Latn":
             translated_answer = await perform_internal_translation(
                 sentences=[answer],
@@ -724,7 +731,6 @@ async def chat_v2(
     logger.info(f"Received prompt: {prompt}, src_lang: {src_lang}, tgt_lang: {tgt_lang}, Image provided: {image is not None}")
 
     try:
-        # Step 1: Handle image if provided
         img = None
         if image:
             image_data = await image.read()
@@ -732,7 +738,6 @@ async def chat_v2(
                 raise HTTPException(status_code=400, detail="Uploaded image is empty")
             img = Image.open(io.BytesIO(image_data))
 
-        # Step 2: Translate prompt to English if needed
         if src_lang != "eng_Latn":
             translated_prompt = await perform_internal_translation(
                 sentences=[prompt],
@@ -745,14 +750,12 @@ async def chat_v2(
             prompt_to_process = prompt
             logger.info("Prompt already in English, no translation needed")
 
-        # Step 3: Generate response in English
         if img:
             response = await llm_manager.chat_v2(img, prompt_to_process)
         else:
             response = await llm_manager.generate(prompt_to_process, settings.max_tokens)
         logger.info(f"Generated English response: {response}")
 
-        # Step 4: Translate response to target language if needed
         if tgt_lang != "eng_Latn":
             translated_response = await perform_internal_translation(
                 sentences=[response],
@@ -797,14 +800,10 @@ async def synthesize_kannada(request: KannadaSynthesizeRequest):
 
 @app.post("/translate", response_model=TranslationResponse)
 async def translate(request: TranslationRequest, translate_manager: TranslateManager = Depends(get_translate_manager)):
-    input_sentences = request.sentences
-    src_lang = request.src_lang
-    tgt_lang = request.tgt_lang
-
-    if not input_sentences:
+    if not request.sentences:
         raise HTTPException(status_code=400, detail="Input sentences are required")
 
-    batch = ip.preprocess_batch(input_sentences, src_lang=src_lang, tgt_lang=tgt_lang)
+    batch = ip.preprocess_batch(request.sentences, src_lang=request.src_lang, tgt_lang=request.tgt_lang)
     inputs = translate_manager.tokenizer(
         batch,
         truncation=True,
@@ -830,25 +829,9 @@ async def translate(request: TranslationRequest, translate_manager: TranslateMan
             clean_up_tokenization_spaces=True,
         )
 
-    translations = ip.postprocess_batch(generated_tokens, lang=tgt_lang)
+    translations = ip.postprocess_batch(generated_tokens, lang=request.tgt_lang)
     return TranslationResponse(translations=translations)
 
-async def perform_internal_translation(sentences: List[str], src_lang: str, tgt_lang: str) -> List[str]:
-    try:
-        translate_manager = model_manager.get_model(src_lang, tgt_lang)
-    except ValueError as e:
-        logger.info(f"Model not preloaded: {str(e)}, loading now...")
-        key = model_manager._get_model_key(src_lang, tgt_lang)
-        model_manager.load_model(src_lang, tgt_lang, key)
-        translate_manager = model_manager.get_model(src_lang, tgt_lang)
-    
-    if not translate_manager.model:
-        translate_manager.load()
-    
-    request = TranslationRequest(sentences=sentences, src_lang=src_lang, tgt_lang=tgt_lang)
-    response = await translate(request, translate_manager)
-    return response.translations
-
 @app.get("/v1/health")
 async def health_check():
     return {"status": "healthy", "model": settings.llm_model_name}