Update app.py

app.py

@@ -1,17 +1,13 @@
-from fastapi import FastAPI, UploadFile, File, Form, HTTPException, BackgroundTasks, WebSocket, Request
+import os
+from fastapi import FastAPI, UploadFile, File, HTTPException
 from fastapi.staticfiles import StaticFiles
 from fastapi.responses import RedirectResponse
-from fastapi_cache import FastAPICache
-from fastapi_cache.backends.redis import RedisBackend
-from fastapi_cache.decorator import cache
 from transformers import pipeline
 import logging
 from PIL import Image
 import io
 from docx import Document
 import fitz  # PyMuPDF
-from pydantic import BaseModel
-import asyncio
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -22,47 +18,80 @@ app = FastAPI()
 # Serve static files (HTML, CSS, JS)
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
-# Load AI models
+# Load a multimodal model for image captioning and visual question answering
 multimodal_pipeline = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
-text_pipeline = pipeline("text2text-generation", model="t5-small")
 
-# Initialize Redis cache
-@app.on_event("startup")
-async def startup():
-    FastAPICache.init(RedisBackend("redis://localhost:6379"))
+# Load a text-based model for summarization and text question answering
+text_pipeline = pipeline("text2text-generation", model="t5-small")
 
-# Root endpoint
 @app.get("/")
 def read_root():
+    # Redirect to the static HTML file
     return RedirectResponse(url="/static/index.html")
 
-# Summarization endpoint
 @app.post("/summarize")
-@cache(expire=300)
-async def summarize_text(file: UploadFile = File(None), text: str = Form(None)):
-    if file:
-        if not file.filename.endswith((".pdf", ".docx")):
-            raise HTTPException(status_code=400, detail="Unsupported file format. Please upload a PDF or DOCX file.")
-        text = await extract_text_from_file(file)
-    elif not text:
-        raise HTTPException(status_code=400, detail="No file or text provided")
-
+async def summarize_text(file: UploadFile = File(...)):
+    logger.info(f"Received document for summarization: {file.filename}")
     try:
+        # Extract text from the document
+        text = await extract_text_from_file(file)
+        
+        # Use the text pipeline to summarize the text
         summary = text_pipeline(f"summarize: {text}", max_length=100)
-        logger.info(f"Generated summary: {summary[0]['summary_text']}")
-        return {"summary": summary[0]['summary_text']}
+        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="An error occurred while processing your request. Please try again.")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/caption")
+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))
+
+        # 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=500, detail=str(e))
+
+@app.post("/answer")
+async def answer_question(file: UploadFile = File(...), question: str = ""):
+    logger.info(f"Received document for question answering: {file.filename}")
+    logger.info(f"Received question: {question}")
+    try:
+        # Extract text from the document
+        text = await extract_text_from_file(file)
+        
+        # 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']}
+    except Exception as e:
+        logger.error(f"Error during question answering: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/vqa")
+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))
 
-# WebSocket for real-time updates
-@app.websocket("/ws")
-async def websocket_endpoint(websocket: WebSocket):
-    await websocket.accept()
-    while True:
-        data = await websocket.receive_text()
-        summary = text_pipeline(f"summarize: {data}", max_length=100)
-        await websocket.send_text(summary[0]['summary_text'])
+        # 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=500, detail=str(e))
 
 # Helper function to extract text from files
 async def extract_text_from_file(file: UploadFile):