Update app.py

app.py

@@ -1,6 +1,8 @@
 """
 A FastAPI application for serving the translation model, inspired by interactive_translate.py.
 """
+
+import os
 import torch
 from transformers import M2M100ForConditionalGeneration, NllbTokenizer
 from fastapi import FastAPI, HTTPException, UploadFile, File
@@ -11,7 +13,10 @@ import logging
 from typing import List
 import fitz  # PyMuPDF
 import shutil
-import os
+
+# ✅ --- 0. Hugging Face Cache Fix ---
+os.environ["HF_HOME"] = "/tmp/huggingface"
+os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface"
 
 # --- 1. App Configuration ---
 logging.basicConfig(level=logging.INFO)
@@ -36,25 +41,30 @@ MODEL_PATH = "facebook/nllb-200-distilled-600M"
 model = None
 tokenizer = None
 
+
 # --- 3. Pydantic Models ---
 class TranslationRequest(BaseModel):
     text: str
     source_language: str
 
+
 class TranslationResponse(BaseModel):
     original_text: str
     translated_text: str
     source_language: str
 
+
 class BatchTranslationRequest(BaseModel):
     texts: List[str]
     source_language: str
 
+
 class BatchTranslationResponse(BaseModel):
     original_texts: List[str]
     translated_texts: List[str]
     source_language: str
-    
+
+
 class PdfTranslationResponse(BaseModel):
     filename: str
     translated_text: str
@@ -72,13 +82,11 @@ def load_model_and_tokenizer(model_path):
         logger.info("Model and tokenizer loaded successfully!")
     except Exception as e:
         logger.error(f"Error loading model: {e}")
-        # In a real app, you might want to exit or handle this more gracefully
         raise
 
+
 def translate_text(text: str, src_lang: str) -> str:
-    """
-    Translates a single string of text to English.
-    """
+    """Translates a single string of text to English."""
     if src_lang not in SUPPORTED_LANGUAGES:
         raise ValueError(f"Language '{src_lang}' not supported.")
 
@@ -93,42 +101,46 @@ def translate_text(text: str, src_lang: str) -> str:
 
     return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
 
+
 def batch_translate_text(texts: List[str], src_lang: str) -> List[str]:
-    """
-    Translates a batch of texts to English.
-    """
+    """Translates a batch of texts to English."""
     if src_lang not in SUPPORTED_LANGUAGES:
         raise ValueError(f"Language '{src_lang}' not supported.")
 
     tokenizer.src_lang = SUPPORTED_LANGUAGES[src_lang]
-    # We use padding=True to handle batches of different lengths
-    inputs = tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=512).to(DEVICE)
+    inputs = tokenizer(
+        texts, return_tensors="pt", padding=True, truncation=True, max_length=512
+    ).to(DEVICE)
 
     generated_tokens = model.generate(
         **inputs,
         forced_bos_token_id=tokenizer.convert_tokens_to_ids("eng_Latn"),
-        max_length=512, # Allow for longer generated sequences in batches
+        max_length=512,
     )
 
     return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
 
+
 # --- 5. API Events ---
 @app.on_event("startup")
 async def startup_event():
     """Load the model at startup."""
     load_model_and_tokenizer(MODEL_PATH)
 
+
 # --- 6. API Endpoints ---
 @app.get("/")
 async def root():
     """Returns the frontend."""
-    return FileResponse('frontend/index.html')
+    return FileResponse("frontend/index.html")
+
 
 @app.get("/languages")
 def get_supported_languages():
     """Returns a list of supported languages."""
     return {"supported_languages": list(SUPPORTED_LANGUAGES.keys())}
 
+
 @app.post("/translate", response_model=TranslationResponse)
 async def translate(request: TranslationRequest):
     """Translates a single text from a source language to English."""
@@ -144,6 +156,7 @@ async def translate(request: TranslationRequest):
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"An unexpected error occurred: {e}")
 
+
 @app.post("/batch-translate", response_model=BatchTranslationResponse)
 async def batch_translate(request: BatchTranslationRequest):
     """Translates a batch of texts from a source language to English."""
@@ -159,19 +172,18 @@ async def batch_translate(request: BatchTranslationRequest):
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"An unexpected error occurred: {e}")
 
+
 @app.post("/translate-pdf", response_model=PdfTranslationResponse)
 async def translate_pdf(source_language: str, file: UploadFile = File(...)):
     """Translates a PDF file from a source language to English."""
     if file.content_type != "application/pdf":
         raise HTTPException(status_code=400, detail="Invalid file type. Please upload a PDF.")
 
-    # Save the uploaded file temporarily
     temp_pdf_path = f"temp_{file.filename}"
     with open(temp_pdf_path, "wb") as buffer:
         shutil.copyfileobj(file.file, buffer)
 
     try:
-        # Extract text from the PDF
         doc = fitz.open(temp_pdf_path)
         extracted_text = ""
         for page in doc:
@@ -181,13 +193,8 @@ async def translate_pdf(source_language: str, file: UploadFile = File(...)):
         if not extracted_text.strip():
             raise HTTPException(status_code=400, detail="Could not extract any text from the PDF.")
 
-        # Split text into chunks (e.g., by paragraph) to handle large texts
-        text_chunks = [p.strip() for p in extracted_text.split('\n') if p.strip()]
-        
-        # Translate the chunks in batches
+        text_chunks = [p.strip() for p in extracted_text.split("\n") if p.strip()]
         translated_chunks = batch_translate_text(text_chunks, source_language)
-
-        # Join the translated chunks back together
         final_translation = "\n".join(translated_chunks)
 
         return PdfTranslationResponse(
@@ -199,15 +206,12 @@ async def translate_pdf(source_language: str, file: UploadFile = File(...)):
         logger.error(f"Error processing PDF: {e}")
         raise HTTPException(status_code=500, detail=f"An error occurred while processing the PDF: {e}")
     finally:
-        # Clean up the temporary file
         if os.path.exists(temp_pdf_path):
             os.remove(temp_pdf_path)
 
 
-# --- 7. Example Usage (for running with uvicorn) ---
-# To run this API, use the following command in your terminal:
-# uvicorn fast_api:app --reload
-
+# --- 7. Run Locally ---
 if __name__ == "__main__":
     import uvicorn
+
     uvicorn.run(app, host="0.0.0.0", port=8000)