add- llm optimisation

Dockerfile

@@ -6,6 +6,6 @@ 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"]

requirements.txt

@@ -1,4 +1,5 @@
 torch
+torchvision
 accelerate
 bitsandbytes
 pillow
@@ -175,7 +176,7 @@ torch==2.6.0
 torchaudio==2.6.0
 torchdiffeq==0.2.5
 tqdm==4.67.1
-transformers==4.50.3
+transformers
 transformers-stream-generator==0.0.5
 triton==3.2.0
 typer==0.15.2

src/server/main.py

@@ -5,7 +5,7 @@ from time import time
 from typing import List
 import tempfile
 import uvicorn
-from fastapi import Depends, FastAPI, File, HTTPException, Query, Request, UploadFile, Body, Form
+from fastapi import Depends, FastAPI, File, HTTPException, Query, Request, UploadFile, Body, Form, APIRouter
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse, RedirectResponse, StreamingResponse
 from PIL import Image
@@ -593,91 +593,10 @@ async def add_request_timing(request: Request, call_next):
 limiter = Limiter(key_func=get_remote_address)
 app.state.limiter = limiter
 
-# API Endpoints
-@app.post("/audio/speech", response_class=StreamingResponse)
-async def synthesize_kannada(request: KannadaSynthesizeRequest):
-    if not tts_manager.model:
-        raise HTTPException(status_code=503, detail="TTS model not loaded")
-    kannada_example = next(ex for ex in EXAMPLES if ex["audio_name"] == "KAN_F (Happy)")
-    if not request.text.strip():
-        raise HTTPException(status_code=400, detail="Text to synthesize cannot be empty.")
-    
-    audio_buffer = synthesize_speech(
-        tts_manager,
-        text=request.text,
-        ref_audio_name="KAN_F (Happy)",
-        ref_text=kannada_example["ref_text"]
-    )
-    
-    return StreamingResponse(
-        audio_buffer,
-        media_type="audio/wav",
-        headers={"Content-Disposition": "attachment; filename=synthesized_kannada_speech.wav"}
-    )
-
-@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:
-        raise HTTPException(status_code=400, detail="Input sentences are required")
+# LLM Router
+llm_router = APIRouter(prefix="/v1", tags=["LLM"])
 
-    batch = ip.preprocess_batch(input_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 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}
-
-@app.get("/")
-async def home():
-    return RedirectResponse(url="/docs")
-
-@app.post("/v1/unload_all_models")
+@llm_router.post("/unload_all_models")
 async def unload_all_models():
     try:
         logger.info("Starting to unload all models...")
@@ -688,7 +607,7 @@ async def unload_all_models():
         logger.error(f"Error unloading models: {str(e)}")
         raise HTTPException(status_code=500, detail=f"Failed to unload models: {str(e)}")
 
-@app.post("/v1/load_all_models")
+@llm_router.post("/load_all_models")
 async def load_all_models():
     try:
         logger.info("Starting to load all models...")
@@ -699,32 +618,15 @@ async def load_all_models():
         logger.error(f"Error loading 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):
-    logger.info(f"Received translation request: {request.dict()}")
-    try:
-        translations = await perform_internal_translation(
-            sentences=request.sentences,
-            src_lang=request.src_lang,
-            tgt_lang=request.tgt_lang
-        )
-        logger.info(f"Translation successful: {translations}")
-        return TranslationResponse(translations=translations)
-    except Exception as e:
-        logger.error(f"Unexpected error during translation: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"Translation failed: {str(e)}")
-
-@app.post("/v1/chat", response_model=ChatResponse)
+@llm_router.post("/chat", response_model=ChatResponse)
 @limiter.limit(settings.chat_rate_limit)
 async def chat(request: Request, chat_request: ChatRequest):
     if not chat_request.prompt:
         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}")
-    
-    EUROPEAN_LANGUAGES = {"deu_Latn", "fra_Latn", "nld_Latn", "spa_Latn", "ita_Latn", "por_Latn", "rus_Cyrl", "pol_Latn"}
-    
     try:
-        if chat_request.src_lang != "eng_Latn" and chat_request.src_lang not in EUROPEAN_LANGUAGES:
+        # 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],
                 src_lang=chat_request.src_lang,
@@ -734,12 +636,14 @@ async def chat(request: Request, chat_request: ChatRequest):
             logger.info(f"Translated prompt to English: {prompt_to_process}")
         else:
             prompt_to_process = chat_request.prompt
-            logger.info("Prompt in English or European language, no translation needed")
+            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 response: {response}")
+        logger.info(f"Generated English response: {response}")
 
-        if chat_request.tgt_lang != "eng_Latn" and chat_request.tgt_lang not in EUROPEAN_LANGUAGES:
+        # Step 3: Translate response to target language if needed
+        if chat_request.tgt_lang != "eng_Latn":
             translated_response = await perform_internal_translation(
                 sentences=[response],
                 src_lang="eng_Latn",
@@ -749,14 +653,14 @@ async def chat(request: Request, chat_request: ChatRequest):
             logger.info(f"Translated response to {chat_request.tgt_lang}: {final_response}")
         else:
             final_response = response
-            logger.info(f"Response in {chat_request.tgt_lang}, no translation needed")
+            logger.info("Response kept in English, no translation needed")
 
         return ChatResponse(response=final_response)
     except Exception as e:
         logger.error(f"Error processing request: {str(e)}")
         raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")
 
-@app.post("/v1/visual_query/")
+@llm_router.post("/visual_query/")
 async def visual_query(
     file: UploadFile = File(...),
     query: str = Body(...),
@@ -768,6 +672,7 @@ 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],
@@ -780,9 +685,11 @@ 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],
@@ -800,7 +707,7 @@ async def visual_query(
         logger.error(f"Error processing request: {str(e)}")
         raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")
 
-@app.post("/v1/chat_v2", response_model=ChatResponse)
+@llm_router.post("/chat_v2", response_model=ChatResponse)
 @limiter.limit(settings.chat_rate_limit)
 async def chat_v2(
     request: Request,
@@ -817,71 +724,154 @@ 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()
             if not image_data:
                 raise HTTPException(status_code=400, detail="Uploaded image is empty")
             img = Image.open(io.BytesIO(image_data))
 
-            if src_lang != "eng_Latn":
-                translated_prompt = await perform_internal_translation(
-                    sentences=[prompt],
-                    src_lang=src_lang,
-                    tgt_lang="eng_Latn"
-                )
-                prompt_to_process = translated_prompt[0]
-                logger.info(f"Translated prompt to English: {prompt_to_process}")
-            else:
-                prompt_to_process = prompt
-                logger.info("Prompt already in English, no translation needed")
-
-            decoded = await llm_manager.chat_v2(img, prompt_to_process)
-            logger.info(f"Generated English response: {decoded}")
-
-            if tgt_lang != "eng_Latn":
-                translated_response = await perform_internal_translation(
-                    sentences=[decoded],
-                    src_lang="eng_Latn",
-                    tgt_lang=tgt_lang
-                )
-                final_response = translated_response[0]
-                logger.info(f"Translated response to {tgt_lang}: {final_response}")
-            else:
-                final_response = decoded
-                logger.info("Response kept in English, no translation needed")
+        # Step 2: Translate prompt to English if needed
+        if src_lang != "eng_Latn":
+            translated_prompt = await perform_internal_translation(
+                sentences=[prompt],
+                src_lang=src_lang,
+                tgt_lang="eng_Latn"
+            )
+            prompt_to_process = translated_prompt[0]
+            logger.info(f"Translated prompt to English: {prompt_to_process}")
         else:
-            if src_lang != "eng_Latn":
-                translated_prompt = await perform_internal_translation(
-                    sentences=[prompt],
-                    src_lang=src_lang,
-                    tgt_lang="eng_Latn"
-                )
-                prompt_to_process = translated_prompt[0]
-                logger.info(f"Translated prompt to English: {prompt_to_process}")
-            else:
-                prompt_to_process = prompt
-                logger.info("Prompt already in English, no translation needed")
+            prompt_to_process = prompt
+            logger.info("Prompt already in English, no translation needed")
 
-            decoded = await llm_manager.generate(prompt_to_process, settings.max_tokens)
-            logger.info(f"Generated English response: {decoded}")
+        # 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}")
 
-            if tgt_lang != "eng_Latn":
-                translated_response = await perform_internal_translation(
-                    sentences=[decoded],
-                    src_lang="eng_Latn",
-                    tgt_lang=tgt_lang
-                )
-                final_response = translated_response[0]
-                logger.info(f"Translated response to {tgt_lang}: {final_response}")
-            else:
-                final_response = decoded
-                logger.info("Response kept in English, no translation needed")
+        # Step 4: Translate response to target language if needed
+        if tgt_lang != "eng_Latn":
+            translated_response = await perform_internal_translation(
+                sentences=[response],
+                src_lang="eng_Latn",
+                tgt_lang=tgt_lang
+            )
+            final_response = translated_response[0]
+            logger.info(f"Translated response to {tgt_lang}: {final_response}")
+        else:
+            final_response = response
+            logger.info("Response kept in English, no translation needed")
 
         return ChatResponse(response=final_response)
     except Exception as e:
         logger.error(f"Error processing request: {str(e)}")
         raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")
 
+# Include LLM Router
+app.include_router(llm_router)
+
+# Other API Endpoints
+@app.post("/audio/speech", response_class=StreamingResponse)
+async def synthesize_kannada(request: KannadaSynthesizeRequest):
+    if not tts_manager.model:
+        raise HTTPException(status_code=503, detail="TTS model not loaded")
+    kannada_example = next(ex for ex in EXAMPLES if ex["audio_name"] == "KAN_F (Happy)")
+    if not request.text.strip():
+        raise HTTPException(status_code=400, detail="Text to synthesize cannot be empty.")
+    
+    audio_buffer = synthesize_speech(
+        tts_manager,
+        text=request.text,
+        ref_audio_name="KAN_F (Happy)",
+        ref_text=kannada_example["ref_text"]
+    )
+    
+    return StreamingResponse(
+        audio_buffer,
+        media_type="audio/wav",
+        headers={"Content-Disposition": "attachment; filename=synthesized_kannada_speech.wav"}
+    )
+
+@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:
+        raise HTTPException(status_code=400, detail="Input sentences are required")
+
+    batch = ip.preprocess_batch(input_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 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}
+
+@app.get("/")
+async def home():
+    return RedirectResponse(url="/docs")
+
+@app.post("/v1/translate", response_model=TranslationResponse)
+async def translate_endpoint(request: TranslationRequest):
+    logger.info(f"Received translation request: {request.dict()}")
+    try:
+        translations = await perform_internal_translation(
+            sentences=request.sentences,
+            src_lang=request.src_lang,
+            tgt_lang=request.tgt_lang
+        )
+        logger.info(f"Translation successful: {translations}")
+        return TranslationResponse(translations=translations)
+    except Exception as e:
+        logger.error(f"Unexpected error during translation: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Translation failed: {str(e)}")
+
 @app.post("/transcribe/", response_model=TranscriptionResponse)
 async def transcribe_audio(file: UploadFile = File(...), language: str = Query(..., enum=list(asr_manager.model_language.keys()))):
     if not asr_manager.model: