test

src/server/main.py

@@ -14,7 +14,7 @@ from pydantic_settings import BaseSettings
 from slowapi import Limiter
 from slowapi.util import get_remote_address
 import torch
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, AutoModel, AutoProcessor, BitsAndBytesConfig
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, AutoProcessor, BitsAndBytesConfig, Gemma3ForConditionalGeneration
 from IndicTransToolkit import IndicProcessor
 import json
 import asyncio
@@ -68,7 +68,7 @@ class Settings(BaseSettings):
 
 settings = Settings()
 
-# Quantization config for LLM
+# Quantization config for LLM (unchanged from gemma_llm.py)
 quantization_config = BitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_quant_type="nf4",
@@ -76,7 +76,7 @@ quantization_config = BitsAndBytesConfig(
     bnb_4bit_compute_dtype=torch.bfloat16
 )
 
-# LLM Manager (adapted from gemma_llm.py)
+# LLM Manager (from gemma_llm.py with async load)
 class LLMManager:
     def __init__(self, model_name: str, device: str = "cuda" if torch.cuda.is_available() else "cpu"):
         self.model_name = model_name
@@ -87,24 +87,11 @@ class LLMManager:
         self.processor = None
         logger.info(f"LLMManager initialized with model {model_name} on {self.device}")
 
-    async def unload(self):
-        if self.is_loaded:
-            await asyncio.to_thread(self._unload_sync)
-            self.is_loaded = False
-            logger.info(f"LLM {self.model_name} unloaded from {self.device}")
-
-    def _unload_sync(self):
-        del self.model
-        del self.processor
-        if self.device.type == "cuda":
-            torch.cuda.empty_cache()
-            logger.info(f"GPU memory allocated after unload: {torch.cuda.memory_allocated()}")
-
     async def load(self):
         if not self.is_loaded:
             try:
                 self.model = await asyncio.to_thread(
-                    AutoModel.from_pretrained,
+                    Gemma3ForConditionalGeneration.from_pretrained,
                     self.model_name,
                     device_map="auto",
                     quantization_config=quantization_config,
@@ -118,6 +105,16 @@ class LLMManager:
                 logger.error(f"Failed to load model: {str(e)}")
                 raise HTTPException(status_code=500, detail=f"Model loading failed: {str(e)}")
 
+    def unload(self):
+        if self.is_loaded:
+            del self.model
+            del self.processor
+            if self.device.type == "cuda":
+                torch.cuda.empty_cache()
+                logger.info(f"GPU memory allocated after unload: {torch.cuda.memory_allocated()}")
+            self.is_loaded = False
+            logger.info(f"LLM {self.model_name} unloaded from {self.device}")
+
     async def generate(self, prompt: str, max_tokens: int = 512, temperature: float = 0.7) -> str:
         if not self.is_loaded:
             await self.load()
@@ -134,15 +131,13 @@ class LLMManager:
         ]
 
         try:
-            inputs_vlm = await asyncio.to_thread(
-                self.processor.apply_chat_template,
+            inputs_vlm = self.processor.apply_chat_template(
                 messages_vlm,
                 add_generation_prompt=True,
                 tokenize=True,
                 return_dict=True,
                 return_tensors="pt"
-            )
-            inputs_vlm = inputs_vlm.to(self.device, dtype=torch.bfloat16)
+            ).to(self.device, dtype=torch.bfloat16)
             logger.info(f"Input IDs: {inputs_vlm['input_ids']}")
             logger.info(f"Decoded input: {self.processor.decode(inputs_vlm['input_ids'][0])}")
         except Exception as e:
@@ -152,8 +147,7 @@ class LLMManager:
         input_len = inputs_vlm["input_ids"].shape[-1]
 
         with torch.inference_mode():
-            generation = await asyncio.to_thread(
-                self.model.generate,
+            generation = self.model.generate(
                 **inputs_vlm,
                 max_new_tokens=max_tokens,
                 do_sample=True,
@@ -188,15 +182,13 @@ class LLMManager:
             logger.info("No valid image provided, processing text only")
 
         try:
-            inputs_vlm = await asyncio.to_thread(
-                self.processor.apply_chat_template,
+            inputs_vlm = self.processor.apply_chat_template(
                 messages_vlm,
                 add_generation_prompt=True,
                 tokenize=True,
                 return_dict=True,
                 return_tensors="pt"
-            )
-            inputs_vlm = inputs_vlm.to(self.device, dtype=torch.bfloat16)
+            ).to(self.device, dtype=torch.bfloat16)
             logger.info(f"Input IDs: {inputs_vlm['input_ids']}")
         except Exception as e:
             logger.error(f"Error in apply_chat_template: {str(e)}")
@@ -205,8 +197,7 @@ class LLMManager:
         input_len = inputs_vlm["input_ids"].shape[-1]
 
         with torch.inference_mode():
-            generation = await asyncio.to_thread(
-                self.model.generate,
+            generation = self.model.generate(
                 **inputs_vlm,
                 max_new_tokens=512,
                 do_sample=True,
@@ -241,15 +232,13 @@ class LLMManager:
             logger.info("No valid image provided, processing text only")
 
         try:
-            inputs_vlm = await asyncio.to_thread(
-                self.processor.apply_chat_template,
+            inputs_vlm = self.processor.apply_chat_template(
                 messages_vlm,
                 add_generation_prompt=True,
                 tokenize=True,
                 return_dict=True,
                 return_tensors="pt"
-            )
-            inputs_vlm = inputs_vlm.to(self.device, dtype=torch.bfloat16)
+            ).to(self.device, dtype=torch.bfloat16)
             logger.info(f"Input IDs: {inputs_vlm['input_ids']}")
         except Exception as e:
             logger.error(f"Error in apply_chat_template: {str(e)}")
@@ -258,8 +247,7 @@ class LLMManager:
         input_len = inputs_vlm["input_ids"].shape[-1]
 
         with torch.inference_mode():
-            generation = await asyncio.to_thread(
-                self.model.generate,
+            generation = self.model.generate(
                 **inputs_vlm,
                 max_new_tokens=512,
                 do_sample=True,
@@ -271,7 +259,7 @@ class LLMManager:
         logger.info(f"Chat_v2 response: {decoded}")
         return decoded
 
-# TTS Manager
+# TTS Manager (async load)
 class TTSManager:
     def __init__(self, device_type=device):
         self.device_type = device_type
@@ -360,32 +348,38 @@ SUPPORTED_LANGUAGES = {
     "por_Latn", "rus_Cyrl", "pol_Latn"
 }
 
-# Translation Manager
+# Translation Manager (async load)
 class TranslateManager:
     def __init__(self, src_lang, tgt_lang, device_type=device, use_distilled=True):
         self.device_type = device_type
-        self.tokenizer, self.model = self.initialize_model(src_lang, tgt_lang, use_distilled)
+        self.tokenizer = None
+        self.model = None
+        self.src_lang = src_lang
+        self.tgt_lang = tgt_lang
+        self.use_distilled = use_distilled
 
-    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.")
-
-        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)
-        model = torch.compile(model, mode="reduce-overhead")
-        print("Model compiled with torch.compile")
-        return tokenizer, model
+    async 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: English to English translation is not supported.")
+
+            self.tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+            self.model = await asyncio.to_thread(
+                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 for {self.src_lang} -> {self.tgt_lang}")
 
 class ModelManager:
     def __init__(self, device_type=device, use_distilled=True, is_lazy_loading=False):
@@ -396,24 +390,9 @@ class ModelManager:
 
     async def load_model(self, src_lang, tgt_lang, key):
         logger.info(f"Loading translation model for {src_lang} -> {tgt_lang}")
-        if src_lang.startswith("eng") and not tgt_lang.startswith("eng"):
-            model_name = "ai4bharat/indictrans2-en-indic-dist-200M" if self.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"
-        else:
-            model_name = "ai4bharat/indictrans2-indic-indic-dist-320M" if self.use_distilled else "ai4bharat/indictrans2-indic-indic-1B"
-        
-        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-        model = await asyncio.to_thread(
-            AutoModelForSeq2SeqLM.from_pretrained,
-            model_name,
-            trust_remote_code=True,
-            torch_dtype=torch.float16,
-            attn_implementation="flash_attention_2"
-        )
-        model = model.to(self.device_type)
-        model = torch.compile(model, mode="reduce-overhead")
-        self.models[key] = TranslateManager(src_lang, tgt_lang, self.device_type, self.use_distilled)
+        translate_manager = TranslateManager(src_lang, tgt_lang, self.device_type, self.use_distilled)
+        await translate_manager.load()
+        self.models[key] = translate_manager
         logger.info(f"Loaded translation model for {key}")
 
     def get_model(self, src_lang, tgt_lang):
@@ -434,7 +413,7 @@ class ModelManager:
             return 'indic_indic'
         raise ValueError("Invalid language combination")
 
-# ASR Manager
+# ASR Manager (async load)
 class ASRModelManager:
     def __init__(self, device_type="cuda"):
         self.device_type = device_type
@@ -483,12 +462,12 @@ class TranslationRequest(BaseModel):
     src_lang: str
     tgt_lang: str
 
-class TranslationResponse(BaseModel):
-    translations: List[str]
-
 class TranscriptionResponse(BaseModel):
     text: str
 
+class TranslationResponse(BaseModel):
+    translations: List[str]
+
 # Dependency
 def get_translate_manager(src_lang: str, tgt_lang: str) -> TranslateManager:
     return model_manager.get_model(src_lang, tgt_lang)
@@ -519,7 +498,7 @@ async def lifespan(app: FastAPI):
     logger.info("Starting model loading in background...")
     asyncio.create_task(load_all_models())
     yield
-    await llm_manager.unload()
+    llm_manager.unload()  # Synchronous unload as per original gemma_llm.py
     logger.info("Server shutdown complete")
 
 # FastAPI App
@@ -604,6 +583,8 @@ async def translate(request: TranslationRequest, translate_manager: TranslateMan
 
 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 translate_manager.model:  # Ensure model is loaded
+        await translate_manager.load()
     request = TranslationRequest(sentences=sentences, src_lang=src_lang, tgt_lang=tgt_lang)
     response = await translate(request, translate_manager)
     return response.translations
@@ -620,7 +601,7 @@ async def home():
 async def unload_all_models():
     try:
         logger.info("Starting to unload all models...")
-        await llm_manager.unload()
+        llm_manager.unload()  # Synchronous as per original
         logger.info("All models unloaded successfully")
         return {"status": "success", "message": "All models unloaded"}
     except Exception as e:
@@ -631,12 +612,12 @@ async def unload_all_models():
 async def load_all_models():
     try:
         logger.info("Starting to load all models...")
-        await llm_manager.load()
+        await llm_manager.load()  # Async load
         logger.info("All models loaded successfully")
         return {"status": "success", "message": "All models loaded"}
     except Exception as e:
         logger.error(f"Error loading models: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"Failed to unload models: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Failed to load models: {str(e)}")
 
 @app.post("/v1/translate", response_model=TranslationResponse)
 async def translate_endpoint(request: TranslationRequest):
@@ -826,7 +807,7 @@ async def transcribe_audio(file: UploadFile = File(...), language: str = Query(.
     if not asr_manager.model:
         raise HTTPException(status_code=503, detail="ASR model still loading, please try again later")
     try:
-        import torchaudio  # Added here for clarity
+        import torchaudio
         wav, sr = torchaudio.load(file.file)
         wav = torch.mean(wav, dim=0, keepdim=True)
         target_sample_rate = 16000