Update app.py

app.py

@@ -1,15 +1,15 @@
 from fastapi import FastAPI, UploadFile, File, HTTPException, Form
 from fastapi.staticfiles import StaticFiles
 from fastapi.responses import RedirectResponse, JSONResponse
-from transformers import pipeline
+from transformers import pipeline, MarianMTModel, MarianTokenizer
 import logging
 from PIL import Image
 import io
 from docx import Document
 import fitz  # PyMuPDF
 import pandas as pd
-import uuid
-from transformers import MarianMTModel, MarianTokenizer
+from tenacity import retry, stop_after_attempt, wait_exponential
+from functools import lru_cache
 from fastapi.middleware.cors import CORSMiddleware
 
 # Configure logging
@@ -30,72 +30,121 @@ app.add_middleware(
 # Serve static files (HTML, CSS, JS)
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
-# Load models
-multimodal_pipeline = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base", use_fast=True)
-text_pipeline = pipeline("text2text-generation", model="t5-small", use_fast=True)
+# Translation models
 translation_models = {
     "fr": "Helsinki-NLP/opus-mt-en-fr",
     "es": "Helsinki-NLP/opus-mt-en-es",
     "de": "Helsinki-NLP/opus-mt-en-de"
 }
 
+# Retry logic for model loading
+@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
+def load_model_with_retry(model_name, task, use_fast=True):
+    logger.info(f"Loading model: {model_name}")
+    return pipeline(task, model=model_name, use_fast=use_fast)
+
+# Lazy-loading pipelines
+@lru_cache(maxsize=1)
+def get_multimodal_pipeline():
+    return load_model_with_retry("Salesforce/blip-image-captioning-base", "image-to-text")
+
+@lru_cache(maxsize=1)
+def get_text_pipeline():
+    return load_model_with_retry("t5-small", "text2text-generation")
+
+@lru_cache(maxsize=3)
+def get_translation_pipeline(target_language):
+    model_name = translation_models.get(target_language, "Helsinki-NLP/opus-mt-en-de")
+    tokenizer = MarianTokenizer.from_pretrained(model_name)
+    model = MarianMTModel.from_pretrained(model_name)
+    return pipeline("translation_en_to_xx", model=model, tokenizer=tokenizer)
+
+# Root endpoint
 @app.get("/")
 def read_root():
     return RedirectResponse(url="/static/index.html")
 
+# Summarize text endpoint
 @app.post("/summarize")
 async def summarize_text(file: UploadFile = File(None), text: str = Form(None)):
-    if file:
-        text = await extract_text_from_file(file)
-    elif not text:
-        raise HTTPException(status_code=400, detail="No text or files provided")
-
-    summary = text_pipeline(f"summarize: {text}", max_length=100)
-    return {"summary": summary[0]['generated_text']}
-
+    try:
+        if file:
+            text = await extract_text_from_file(file)
+        elif not text:
+            raise HTTPException(status_code=400, detail="No text or file provided")
+
+        text_pipeline = get_text_pipeline()
+        summary = text_pipeline(f"summarize: {text}", max_length=100)
+        return {"summary": summary[0]['generated_text']}
+    except Exception as e:
+        logger.error(f"Error in summarization: {e}")
+        raise HTTPException(status_code=500, detail="Failed to summarize text. Please try again.")
+
+# Image captioning endpoint
 @app.post("/caption")
 async def caption_image(file: UploadFile = File(...)):
-    image_data = await file.read()
-    image = Image.open(io.BytesIO(image_data))
-    caption = multimodal_pipeline(image)
-    return {"caption": caption[0]['generated_text']}
-
+    try:
+        image_data = await file.read()
+        image = Image.open(io.BytesIO(image_data))
+        multimodal_pipeline = get_multimodal_pipeline()
+        caption = multimodal_pipeline(image)
+        return {"caption": caption[0]['generated_text']}
+    except Exception as e:
+        logger.error(f"Error in image captioning: {e}")
+        raise HTTPException(status_code=500, detail="Failed to generate caption. Please try again.")
+
+# Translation endpoint
 @app.post("/translate")
 async def translate_document(file: UploadFile = File(None), text: str = Form(None), target_language: str = Form(...)):
-    if file:
-        text = await extract_text_from_file(file)
-    elif not text:
-        raise HTTPException(status_code=400, detail="No text or file provided")
-
-    model_name = translation_models.get(target_language, "Helsinki-NLP/opus-mt-en-de")
-    tokenizer = MarianTokenizer.from_pretrained(model_name)
-    model = MarianMTModel.from_pretrained(model_name)
-    translated = model.generate(**tokenizer(text, return_tensors="pt", truncation=True))
-    translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
-    return {"translated_text": translated_text}
-
+    try:
+        if file:
+            text = await extract_text_from_file(file)
+        elif not text:
+            raise HTTPException(status_code=400, detail="No text or file provided")
+
+        translation_pipeline = get_translation_pipeline(target_language)
+        translated = translation_pipeline(text)
+        return {"translated_text": translated[0]['translation_text']}
+    except Exception as e:
+        logger.error(f"Error in translation: {e}")
+        raise HTTPException(status_code=500, detail="Failed to translate text. Please try again.")
+
+# Question answering endpoint
 @app.post("/answer")
 async def answer_question(file: UploadFile = File(None), text: str = Form(None), question: str = Form(...)):
-    if file:
-        text = await extract_text_from_file(file)
-    elif not text:
-        raise HTTPException(status_code=400, detail="No text or file provided")
-
-    answer = text_pipeline(f"question: {question} context: {text}")
-    return {"answer": answer[0]['generated_text']}
-
+    try:
+        if file:
+            text = await extract_text_from_file(file)
+        elif not text:
+            raise HTTPException(status_code=400, detail="No text or file provided")
+
+        text_pipeline = get_text_pipeline()
+        answer = text_pipeline(f"question: {question} context: {text}")
+        return {"answer": answer[0]['generated_text']}
+    except Exception as e:
+        logger.error(f"Error in question answering: {e}")
+        raise HTTPException(status_code=500, detail="Failed to answer the question. Please try again.")
+
+# Visual question answering endpoint
 @app.post("/vqa")
 async def visual_question_answering(file: UploadFile = File(...), question: str = Form(...)):
-    image_data = await file.read()
-    image = Image.open(io.BytesIO(image_data))
-    answer = multimodal_pipeline(image, question=question)
-    return {"answer": answer[0]['generated_text']}
-
+    try:
+        image_data = await file.read()
+        image = Image.open(io.BytesIO(image_data))
+        multimodal_pipeline = get_multimodal_pipeline()
+        answer = multimodal_pipeline(image, question=question)
+        return {"answer": answer[0]['generated_text']}
+    except Exception as e:
+        logger.error(f"Error in visual question answering: {e}")
+        raise HTTPException(status_code=500, detail="Failed to answer the question. Please try again.")
+
+# Data visualization endpoint
 @app.post("/visualize")
 async def visualize_data(file: UploadFile = File(...), request: str = Form(...)):
-    df = pd.read_excel(io.BytesIO(await file.read()))
-    if "bar" in request.lower():
-        code = f"""
+    try:
+        df = pd.read_excel(io.BytesIO(await file.read()))
+        if "bar" in request.lower():
+            code = f"""
 import matplotlib.pyplot as plt
 plt.bar(df['{df.columns[0]}'], df['{df.columns[1]}'])
 plt.xlabel('{df.columns[0]}')
@@ -103,8 +152,8 @@ plt.ylabel('{df.columns[1]}')
 plt.title('Bar Chart')
 plt.show()
 """
-    elif "line" in request.lower():
-        code = f"""
+        elif "line" in request.lower():
+            code = f"""
 import matplotlib.pyplot as plt
 plt.plot(df['{df.columns[0]}'], df['{df.columns[1]}'])
 plt.xlabel('{df.columns[0]}')
@@ -112,30 +161,39 @@ plt.ylabel('{df.columns[1]}')
 plt.title('Line Chart')
 plt.show()
 """
-    else:
-        code = f"""
+        else:
+            code = f"""
 import seaborn as sns
 sns.pairplot(df)
 plt.show()
 """
-    return {"code": code}
+        return {"code": code}
+    except Exception as e:
+        logger.error(f"Error in data visualization: {e}")
+        raise HTTPException(status_code=500, detail="Failed to generate visualization code. Please try again.")
 
+# Helper function to extract text from files
 async def extract_text_from_file(file: UploadFile):
-    file_content = await file.read()
-    if file.filename.endswith(".pdf"):
-        doc = fitz.open(stream=file_content, filetype="pdf")
-        text = ""
-        for page in doc:
-            text += page.get_text()
-        return text
-    elif file.filename.endswith(".docx"):
-        doc = Document(io.BytesIO(file_content))
-        return "\n".join([para.text for para in doc.paragraphs])
-    elif file.filename.endswith(".txt"):
-        return file_content.decode("utf-8")
-    else:
-        raise HTTPException(status_code=400, detail="Unsupported file format")
-
+    try:
+        file_content = await file.read()
+        if file.filename.endswith(".pdf"):
+            doc = fitz.open(stream=file_content, filetype="pdf")
+            text = ""
+            for page in doc:
+                text += page.get_text()
+            return text
+        elif file.filename.endswith(".docx"):
+            doc = Document(io.BytesIO(file_content))
+            return "\n".join([para.text for para in doc.paragraphs])
+        elif file.filename.endswith(".txt"):
+            return file_content.decode("utf-8")
+        else:
+            raise HTTPException(status_code=400, detail="Unsupported file format")
+    except Exception as e:
+        logger.error(f"Error extracting text from file: {e}")
+        raise HTTPException(status_code=500, detail="Failed to extract text from file. Please try again.")
+
+# Run the application
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)