Update app.py

app.py

@@ -1,17 +1,18 @@
+import os
 from fastapi import FastAPI, UploadFile, File, HTTPException, Form
 from fastapi.staticfiles import StaticFiles
-from fastapi.responses import RedirectResponse, JSONResponse
-from transformers import pipeline, MarianMTModel, MarianTokenizer
+from fastapi.responses import RedirectResponse
+from transformers import pipeline
 import logging
 from PIL import Image
 import io
 from docx import Document
 import fitz  # PyMuPDF
 import pandas as pd
-from tenacity import retry, stop_after_attempt, wait_exponential
-from functools import lru_cache
-from fastapi.middleware.cors import CORSMiddleware
-import os
+import matplotlib.pyplot as plt
+import seaborn as sns
+import uuid
+from transformers import MarianMTModel, MarianTokenizer
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -19,76 +20,51 @@ logger = logging.getLogger(__name__)
 
 app = FastAPI()
 
-# Add CORS middleware
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],  # Allow all origins (replace with your frontend URL in production)
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
 # Serve static files (HTML, CSS, JS)
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
-# Translation models
+# 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 = {
     "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)):
-    try:
-        if file:
-            logger.info(f"Received document for summarization: {file.filename}")
+async def summarize_text(
+    file: UploadFile = File(None),
+    text: str = Form(None)
+):
+    if file:
+        logger.info(f"Received document for summarization: {file.filename}")
+        try:
             text = await extract_text_from_file(file)
-        elif text:
-            logger.info("Received manual text for summarization")
-        else:
-            raise HTTPException(status_code=400, detail="No file or text provided")
+        except Exception as e:
+            logger.error(f"Error extracting text from file: {e}")
+            raise HTTPException(status_code=400, detail=str(e))
+    elif text:
+        logger.info("Received manual text for summarization")
+    else:
+        raise HTTPException(status_code=400, detail="No file or text provided")
 
-        text_pipeline = get_text_pipeline()
+    try:
         summary = text_pipeline(f"summarize: {text}", max_length=100)
+        logger.info(f"Generated summary: {summary[0]['generated_text']}")
         return {"summary": summary[0]['generated_text']}
     except Exception as e:
         logger.error(f"Error during summarization: {e}")
-        raise HTTPException(status_code=500, detail="Failed to summarize text. Please try again.")
+        raise HTTPException(status_code=500, detail=str(e))
 
-# Image captioning endpoint
 @app.post("/caption")
 async def caption_image(file: UploadFile = File(...)):
+    logger.info(f"Received image for captioning: {file.filename}")
     try:
-        logger.info(f"Received image for captioning: {file.filename}")
         image_data = await file.read()
         image = Image.open(io.BytesIO(image_data))
 
@@ -96,55 +72,71 @@ async def caption_image(file: UploadFile = File(...)):
         if image.format not in ["JPEG", "PNG"]:
             raise ValueError("Unsupported image format. Please upload a JPEG or PNG file.")
 
-        multimodal_pipeline = get_multimodal_pipeline()
         caption = multimodal_pipeline(image)
+        logger.info(f"Generated caption: {caption[0]['generated_text']}")
         return {"caption": caption[0]['generated_text']}
     except Exception as e:
         logger.error(f"Error during image captioning: {e}")
         raise HTTPException(status_code=400, detail=str(e))
 
-# Translation endpoint
 @app.post("/translate")
-async def translate_document(file: UploadFile = File(None), text: str = Form(None), target_language: str = Form(...)):
+async def translate_document(
+    file: UploadFile = File(...),
+    target_language: str = Form(...)
+):
+    logger.info(f"Received document for translation: {file.filename}")
+    logger.info(f"Target language: {target_language}")
+
     try:
-        if file:
-            logger.info(f"Received document for translation: {file.filename}")
-            text = await extract_text_from_file(file)
-        elif not text:
-            raise HTTPException(status_code=400, detail="No text or file provided")
+        text = await extract_text_from_file(file)
 
-        translation_pipeline = get_translation_pipeline(target_language)
-        translated = translation_pipeline(text)
-        return {"translated_text": translated[0]['translation_text']}
+        if target_language in translation_models:
+            model_name = translation_models[target_language]
+        else:
+            model_name = "Helsinki-NLP/opus-mt-en-de"  # Default to German
+
+        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}
     except Exception as e:
-        logger.error(f"Error during translation: {e}")
-        raise HTTPException(status_code=500, detail="Failed to translate text. Please try again.")
+        logger.error(f"Error during document translation: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
 
-# Question answering endpoint
 @app.post("/answer")
-async def answer_question(file: UploadFile = File(None), text: str = Form(None), question: str = Form(...)):
-    try:
-        if file:
-            logger.info(f"Received document for question answering: {file.filename}")
+async def answer_question(
+    file: UploadFile = File(None),
+    text: str = Form(None),
+    question: str = Form(...)
+):
+    if file:
+        logger.info(f"Received document for question answering: {file.filename}")
+        try:
             text = await extract_text_from_file(file)
-        elif text:
-            logger.info("Received manual text for question answering")
-        else:
-            raise HTTPException(status_code=400, detail="No file or text provided")
+        except Exception as e:
+            logger.error(f"Error extracting text from file: {e}")
+            raise HTTPException(status_code=400, detail=str(e))
+    elif text:
+        logger.info("Received manual text for question answering")
+    else:
+        raise HTTPException(status_code=400, detail="No file or text provided")
 
-        text_pipeline = get_text_pipeline()
+    try:
         answer = text_pipeline(f"question: {question} context: {text}")
+        logger.info(f"Generated answer: {answer[0]['generated_text']}")
         return {"answer": answer[0]['generated_text']}
     except Exception as e:
         logger.error(f"Error during question answering: {e}")
-        raise HTTPException(status_code=500, detail="Failed to answer the question. Please try again.")
+        raise HTTPException(status_code=500, detail=str(e))
 
-# Visual question answering endpoint
 @app.post("/vqa")
 async def visual_question_answering(file: UploadFile = File(...), question: str = Form(...)):
+    logger.info(f"Received image for visual question answering: {file.filename}")
+    logger.info(f"Received question: {question}")
     try:
-        logger.info(f"Received image for visual question answering: {file.filename}")
-        logger.info(f"Received question: {question}")
         image_data = await file.read()
         image = Image.open(io.BytesIO(image_data))
 
@@ -152,19 +144,22 @@ async def visual_question_answering(file: UploadFile = File(...), question: str
         if image.format not in ["JPEG", "PNG"]:
             raise ValueError("Unsupported image format. Please upload a JPEG or PNG file.")
 
-        multimodal_pipeline = get_multimodal_pipeline()
         answer = multimodal_pipeline(image, question=question)
+        logger.info(f"Generated answer: {answer[0]['generated_text']}")
         return {"answer": answer[0]['generated_text']}
     except Exception as e:
         logger.error(f"Error during visual question answering: {e}")
         raise HTTPException(status_code=400, detail=str(e))
 
-# Data visualization endpoint
 @app.post("/visualize")
-async def visualize_data(file: UploadFile = File(...), request: str = Form(...)):
+async def visualize_data(
+    file: UploadFile = File(...),
+    request: str = Form(...)
+):
+    logger.info(f"Received Excel file for visualization: {file.filename}")
+    logger.info(f"Received visualization request: {request}")
+
     try:
-        logger.info(f"Received Excel file for visualization: {file.filename}")
-        logger.info(f"Received visualization request: {request}")
         df = pd.read_excel(io.BytesIO(await file.read()))
 
         if "bar" in request.lower():
@@ -192,12 +187,15 @@ sns.pairplot(df)
 plt.show()
 """
 
-        return {"code": code}
+        code_filename = f"visualization_{uuid.uuid4()}.py"
+        with open(code_filename, "w") as f:
+            f.write(code)
+
+        return {"code": code, "filename": code_filename}
     except Exception as e:
         logger.error(f"Error during visualization code generation: {e}")
-        raise HTTPException(status_code=500, detail="Failed to generate visualization code. Please try again.")
+        raise HTTPException(status_code=500, detail=str(e))
 
-# Helper function to extract text from files
 async def extract_text_from_file(file: UploadFile):
     try:
         file_content = await file.read()
@@ -209,16 +207,16 @@ async def extract_text_from_file(file: UploadFile):
             return text
         elif file.filename.endswith(".docx"):
             doc = Document(io.BytesIO(file_content))
-            return "\n".join([para.text for para in doc.paragraphs])
+            text = "\n".join([para.text for para in doc.paragraphs])
+            return text
         elif file.filename.endswith(".txt"):
             return file_content.decode("utf-8")
         else:
-            raise HTTPException(status_code=400, detail="Unsupported file format")
+            raise ValueError("Unsupported file format. Please upload a PDF, DOCX, or TXT file.")
     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.")
+        raise HTTPException(status_code=400, detail=str(e))
 
-# Run the application
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)