Update app.py

app.py

@@ -23,9 +23,13 @@ app = FastAPI()
 # 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)
+# Load a multimodal model for image captioning and visual question answering
+multimodal_pipeline = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
+
+# Load a text-based model for summarization and text question answering
+text_pipeline = pipeline("text2text-generation", model="t5-small")
+
+# Load a translation model (initialized dynamically based on target language)
 translation_models = {
     "fr": "Helsinki-NLP/opus-mt-en-fr",
     "es": "Helsinki-NLP/opus-mt-en-es",
@@ -34,16 +38,18 @@ translation_models = {
 
 @app.get("/")
 def read_root():
+    # Redirect to the static HTML file
     return RedirectResponse(url="/static/index.html")
 
 @app.post("/summarize")
 async def summarize_text(
-    file: UploadFile = File(None),
-    text: str = Form(None)
+    file: UploadFile = File(None),  # Optional file upload
+    text: str = Form(None)  # Optional manual text input
 ):
     if file:
         logger.info(f"Received document for summarization: {file.filename}")
         try:
+            # Extract text from the document
             text = await extract_text_from_file(file)
         except Exception as e:
             logger.error(f"Error extracting text from file: {e}")
@@ -54,6 +60,7 @@ async def summarize_text(
         raise HTTPException(status_code=400, detail="No file or text provided")
 
     try:
+        # Use the text pipeline to summarize the text
         summary = text_pipeline(f"summarize: {text}", max_length=100)
         logger.info(f"Generated summary: {summary[0]['generated_text']}")
         return {"summary": summary[0]['generated_text']}
@@ -65,56 +72,28 @@ async def summarize_text(
 async def caption_image(file: UploadFile = File(...)):
     logger.info(f"Received image for captioning: {file.filename}")
     try:
+        # Read the image file
         image_data = await file.read()
         image = Image.open(io.BytesIO(image_data))
 
-        # Validate image format
-        if image.format not in ["JPEG", "PNG"]:
-            raise ValueError("Unsupported image format. Please upload a JPEG or PNG file.")
-
+        # Use the multimodal pipeline to generate a caption for the image
         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))
-
-@app.post("/translate")
-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:
-        text = await extract_text_from_file(file)
-
-        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 document translation: {e}")
         raise HTTPException(status_code=500, detail=str(e))
 
 @app.post("/answer")
 async def answer_question(
-    file: UploadFile = File(None),
-    text: str = Form(None),
-    question: str = Form(...)
+    file: UploadFile = File(None),  # Optional file upload
+    text: str = Form(None),  # Optional manual text input
+    question: str = Form(...)  # Required question
 ):
     if file:
         logger.info(f"Received document for question answering: {file.filename}")
         try:
+            # Extract text from the document
             text = await extract_text_from_file(file)
         except Exception as e:
             logger.error(f"Error extracting text from file: {e}")
@@ -125,6 +104,7 @@ async def answer_question(
         raise HTTPException(status_code=400, detail="No file or text provided")
 
     try:
+        # Use the text pipeline to answer the question
         answer = text_pipeline(f"question: {question} context: {text}")
         logger.info(f"Generated answer: {answer[0]['generated_text']}")
         return {"answer": answer[0]['generated_text']}
@@ -137,19 +117,17 @@ async def visual_question_answering(file: UploadFile = File(...), question: str
     logger.info(f"Received image for visual question answering: {file.filename}")
     logger.info(f"Received question: {question}")
     try:
+        # Read the image file
         image_data = await file.read()
         image = Image.open(io.BytesIO(image_data))
 
-        # Validate image format
-        if image.format not in ["JPEG", "PNG"]:
-            raise ValueError("Unsupported image format. Please upload a JPEG or PNG file.")
-
+        # Use the multimodal pipeline to answer the question about the image
         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))
+        raise HTTPException(status_code=500, detail=str(e))
 
 @app.post("/visualize")
 async def visualize_data(
@@ -160,8 +138,10 @@ async def visualize_data(
     logger.info(f"Received visualization request: {request}")
 
     try:
+        # Read the Excel file
         df = pd.read_excel(io.BytesIO(await file.read()))
 
+        # Generate visualization code based on the request
         if "bar" in request.lower():
             code = f"""
 import matplotlib.pyplot as plt
@@ -187,6 +167,7 @@ sns.pairplot(df)
 plt.show()
 """
 
+        # Save the generated code to a file (optional)
         code_filename = f"visualization_{uuid.uuid4()}.py"
         with open(code_filename, "w") as f:
             f.write(code)
@@ -196,27 +177,56 @@ plt.show()
         logger.error(f"Error during visualization code generation: {e}")
         raise HTTPException(status_code=500, detail=str(e))
 
+@app.post("/translate")
+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:
+        # Extract text from the document
+        text = await extract_text_from_file(file)
+
+        # Load a translation model based on the target language
+        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)
+
+        # Translate the text
+        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 document translation: {e}")
+        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()
         if file.filename.endswith(".pdf"):
-            doc = fitz.open(stream=file_content, filetype="pdf")
+            doc = fitz.open(stream=await file.read(), filetype="pdf")
             text = ""
             for page in doc:
                 text += page.get_text()
             return text
         elif file.filename.endswith(".docx"):
-            doc = Document(io.BytesIO(file_content))
+            doc = Document(io.BytesIO(await file.read()))
             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 ValueError("Unsupported file format. Please upload a PDF, DOCX, or TXT file.")
+            raise ValueError("Unsupported file format. Please upload a PDF or DOCX file.")
     except Exception as e:
         logger.error(f"Error extracting text from file: {e}")
         raise HTTPException(status_code=400, detail=str(e))
 
+# Hugging Face Spaces expects the app to be served on port 7860
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)