-co-locate-translatio

Dockerfile

@@ -3,19 +3,19 @@ WORKDIR /app
 
 RUN apt-get update && apt-get install -y \
     python3 \
-    python3-pip \
+    python3-pip python3-distutils python3-dev python3-venv\
     git \
     ffmpeg \
-    sudo \
-    wget libvips\
-    build-essential \          
-    curl \   
+    sudo wget curl software-properties-common build-essential gcc g++ \   
     && ln -s /usr/bin/python3 /usr/bin/python \
     && rm -rf /var/lib/apt/lists/*
 
 RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
 ENV PATH="/root/.cargo/bin:${PATH}"
 
+RUN export CC=/usr/bin/gcc
+RUN export CXX=/usr/bin/g++
+
 RUN pip install --upgrade pip setuptools setuptools-rust torch
 COPY requirements.txt .
 #RUN pip install --no-cache-dir torch==2.6.0 torchvision

requirements.txt

@@ -7,4 +7,5 @@ uvicorn
 fastapi
 pydantic_settings
 slowapi
-python-multipart
+python-multipart
+IndicTransToolkit @ git+https://github.com/VarunGumma/IndicTransToolkit.git@399b3fec93d2ee85cb998cb7a4fb7a7d83afcbcf

src/server/main.py

@@ -14,12 +14,14 @@ from pydantic import BaseModel, field_validator
 from pydantic_settings import BaseSettings
 from slowapi import Limiter
 from slowapi.util import get_remote_address
-import requests
+import torch
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from IndicTransToolkit import IndicProcessor
 
 from logging_config import logger
 from tts_config import SPEED, ResponseFormat, config as tts_config
 from gemma_llm import LLMManager
-#from auth import get_api_key, settings as auth_settings
+# from auth import get_api_key, settings as auth_settings
 
 # Supported language codes
 SUPPORTED_LANGUAGES = {
@@ -68,6 +70,73 @@ app.state.limiter = limiter
 
 llm_manager = LLMManager(settings.llm_model_name)
 
+# Translation Manager and Model Manager
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+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)
+
+    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)
+        return tokenizer, model
+
+class ModelManager:
+    def __init__(self, device_type=DEVICE, use_distilled=True, is_lazy_loading=False):
+        self.models: dict[str, TranslateManager] = {}
+        self.device_type = device_type
+        self.use_distilled = use_distilled
+        self.is_lazy_loading = is_lazy_loading
+        if not is_lazy_loading:
+            self.preload_models()
+
+    def preload_models(self):
+        self.models['eng_indic'] = TranslateManager('eng_Latn', 'kan_Knda', self.device_type, self.use_distilled)
+        self.models['indic_eng'] = TranslateManager('kan_Knda', 'eng_Latn', self.device_type, self.use_distilled)
+        self.models['indic_indic'] = TranslateManager('kan_Knda', 'hin_Deva', self.device_type, self.use_distilled)
+
+    def get_model(self, src_lang, tgt_lang) -> TranslateManager:
+        if src_lang.startswith("eng") and not tgt_lang.startswith("eng"):
+            key = 'eng_indic'
+        elif not src_lang.startswith("eng") and tgt_lang.startswith("eng"):
+            key = 'indic_eng'
+        elif not src_lang.startswith("eng") and not tgt_lang.startswith("eng"):
+            key = 'indic_indic'
+        else:
+            raise ValueError("Invalid language combination: English to English translation is not supported.")
+
+        if key not in self.models:
+            if self.is_lazy_loading:
+                if key == 'eng_indic':
+                    self.models[key] = TranslateManager('eng_Latn', 'kan_Knda', self.device_type, self.use_distilled)
+                elif key == 'indic_eng':
+                    self.models[key] = TranslateManager('kan_Knda', 'eng_Latn', self.device_type, self.use_distilled)
+                elif key == 'indic_indic':
+                    self.models[key] = TranslateManager('kan_Knda', 'hin_Deva', self.device_type, self.use_distilled)
+            else:
+                raise ValueError(f"Model for {key} is not preloaded and lazy loading is disabled.")
+        return self.models[key]
+
+ip = IndicProcessor(inference=True)
+model_manager = ModelManager()
+
+# Pydantic Models
 class ChatRequest(BaseModel):
     prompt: str
     src_lang: str = "kan_Knda"  # Default to Kannada
@@ -93,48 +162,61 @@ class TranslationRequest(BaseModel):
     src_lang: str
     tgt_lang: str
 
-    @field_validator("src_lang", "tgt_lang")
-    def validate_language(cls, v):
-        if v not in SUPPORTED_LANGUAGES:
-            raise ValueError(f"Unsupported language code: {v}. Supported codes: {', '.join(SUPPORTED_LANGUAGES)}")
-        return v
-
 class TranslationResponse(BaseModel):
     translations: List[str]
 
-async def call_external_translation(sentences: List[str], src_lang: str, tgt_lang: str) -> List[str]:
-    external_url = "https://gaganyatri-dhwani-server.hf.space/v1/translate"
-    payload = {
-        "sentences": sentences,
-        "src_lang": src_lang,
-        "tgt_lang": tgt_lang
-    }
-    try:
-        response = requests.post(
-            external_url,
-            json=payload,
-            headers={
-                "accept": "application/json",
-                "Content-Type": "application/json"
-            },
-            timeout=60
+# Dependency to get TranslateManager
+def get_translate_manager(src_lang: str, tgt_lang: str) -> TranslateManager:
+    return model_manager.get_model(src_lang, tgt_lang)
+
+# Internal Translation Endpoint
+@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,
         )
-        response.raise_for_status()
-        translations = response.json().get("translations", [])
-        if not translations or len(translations) != len(sentences):
-            logger.warning(f"Unexpected response format: {response.json()}")
-            raise ValueError("Invalid response from translation service")
-        return translations
-    except requests.Timeout:
-        logger.error("Translation request timed out")
-        raise HTTPException(status_code=504, detail="Translation service timeout")
-    except requests.RequestException as e:
-        logger.error(f"Error during translation: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"Translation failed: {str(e)}")
-    except ValueError as e:
-        logger.error(f"Invalid response: {str(e)}")
-        raise HTTPException(status_code=500, detail=str(e))
 
+    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)
+
+# Helper function to perform internal translation
+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)
+    request = TranslationRequest(sentences=sentences, src_lang=src_lang, tgt_lang=tgt_lang)
+    response = await translate(request, translate_manager)
+    return response.translations
+
+# API Endpoints
 @app.get("/v1/health")
 async def health_check():
     return {"status": "healthy", "model": settings.llm_model_name}
@@ -144,9 +226,7 @@ async def home():
     return RedirectResponse(url="/docs")
 
 @app.post("/v1/unload_all_models")
-async def unload_all_models(
-    #api_key: str = Depends(get_api_key)
-    ):
+async def unload_all_models():
     try:
         logger.info("Starting to unload all models...")
         llm_manager.unload()
@@ -157,9 +237,7 @@ async def unload_all_models(
         raise HTTPException(status_code=500, detail=f"Failed to unload models: {str(e)}")
 
 @app.post("/v1/load_all_models")
-async def load_all_models(
-    #api_key: str = Depends(get_api_key)
-    ):
+async def load_all_models():
     try:
         logger.info("Starting to load all models...")
         llm_manager.load()
@@ -170,34 +248,30 @@ async def load_all_models(
         raise HTTPException(status_code=500, detail=f"Failed to load models: {str(e)}")
 
 @app.post("/v1/translate", response_model=TranslationResponse)
-async def translate(request: TranslationRequest):
+async def translate_endpoint(request: TranslationRequest):
     logger.info(f"Received translation request: {request.dict()}")
     try:
-        translations = await call_external_translation(
+        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 HTTPException as e:
-        raise e
     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)
 @limiter.limit(settings.chat_rate_limit)
-async def chat(request: Request, chat_request: ChatRequest,
-                #api_key: str = Depends(get_api_key)
-                ):
+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}")
     try:
         # Translate prompt to English if src_lang is not English
         if chat_request.src_lang != "eng_Latn":
-            translated_prompt = await call_external_translation(
+            translated_prompt = await perform_internal_translation(
                 sentences=[chat_request.prompt],
                 src_lang=chat_request.src_lang,
                 tgt_lang="eng_Latn"
@@ -214,7 +288,7 @@ async def chat(request: Request, chat_request: ChatRequest,
 
         # Translate response to target language if tgt_lang is not English
         if chat_request.tgt_lang != "eng_Latn":
-            translated_response = await call_external_translation(
+            translated_response = await perform_internal_translation(
                 sentences=[response],
                 src_lang="eng_Latn",
                 tgt_lang=chat_request.tgt_lang
@@ -236,16 +310,15 @@ async def visual_query(
     query: str = Body(...),
     src_lang: str = Query("kan_Knda", enum=list(SUPPORTED_LANGUAGES)),
     tgt_lang: str = Query("kan_Knda", enum=list(SUPPORTED_LANGUAGES)),
-    #api_key: str = Depends(get_api_key)
 ):
     try:
         image = Image.open(file.file)
         if image.size == (0, 0):
             raise HTTPException(status_code=400, detail="Uploaded image is empty or invalid")
-        
+
         # Translate query to English if src_lang is not English
         if src_lang != "eng_Latn":
-            translated_query = await call_external_translation(
+            translated_query = await perform_internal_translation(
                 sentences=[query],
                 src_lang=src_lang,
                 tgt_lang="eng_Latn"
@@ -262,7 +335,7 @@ async def visual_query(
 
         # Translate answer to target language if tgt_lang is not English
         if tgt_lang != "eng_Latn":
-            translated_answer = await call_external_translation(
+            translated_answer = await perform_internal_translation(
                 sentences=[answer],
                 src_lang="eng_Latn",
                 tgt_lang=tgt_lang
@@ -286,13 +359,12 @@ async def chat_v2(
     image: UploadFile = File(default=None),
     src_lang: str = Form("kan_Knda"),
     tgt_lang: str = Form("kan_Knda"),
-    #api_key: str = Depends(get_api_key)
 ):
     if not prompt:
         raise HTTPException(status_code=400, detail="Prompt cannot be empty")
     if src_lang not in SUPPORTED_LANGUAGES or tgt_lang not in SUPPORTED_LANGUAGES:
         raise HTTPException(status_code=400, detail=f"Unsupported language code. Supported codes: {', '.join(SUPPORTED_LANGUAGES)}")
-    
+
     logger.info(f"Received prompt: {prompt}, src_lang: {src_lang}, tgt_lang: {tgt_lang}, Image provided: {image is not None}")
 
     try:
@@ -301,10 +373,10 @@ async def chat_v2(
             if not image_data:
                 raise HTTPException(status_code=400, detail="Uploaded image is empty")
             img = Image.open(io.BytesIO(image_data))
-            
+
             # Translate prompt to English if src_lang is not English
             if src_lang != "eng_Latn":
-                translated_prompt = await call_external_translation(
+                translated_prompt = await perform_internal_translation(
                     sentences=[prompt],
                     src_lang=src_lang,
                     tgt_lang="eng_Latn"
@@ -320,7 +392,7 @@ async def chat_v2(
 
             # Translate response to target language if tgt_lang is not English
             if tgt_lang != "eng_Latn":
-                translated_response = await call_external_translation(
+                translated_response = await perform_internal_translation(
                     sentences=[decoded],
                     src_lang="eng_Latn",
                     tgt_lang=tgt_lang
@@ -333,7 +405,7 @@ async def chat_v2(
         else:
             # Translate prompt to English if src_lang is not English
             if src_lang != "eng_Latn":
-                translated_prompt = await call_external_translation(
+                translated_prompt = await perform_internal_translation(
                     sentences=[prompt],
                     src_lang=src_lang,
                     tgt_lang="eng_Latn"
@@ -343,13 +415,13 @@ async def chat_v2(
             else:
                 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}")
-            
+
             # Translate response to target language if tgt_lang is not English
             if tgt_lang != "eng_Latn":
-                translated_response = await call_external_translation(
+                translated_response = await perform_internal_translation(
                     sentences=[decoded],
                     src_lang="eng_Latn",
                     tgt_lang=tgt_lang