Update app.py

app.py

@@ -1,174 +1,157 @@
 import os
-from fastapi import FastAPI, UploadFile, File, HTTPException, Form
+from fastapi import FastAPI, UploadFile, File, Form, HTTPException
 from fastapi.staticfiles import StaticFiles
-from fastapi.responses import RedirectResponse, JSONResponse
+from fastapi.responses import HTMLResponse, JSONResponse
 from transformers import pipeline
-import logging
 from PIL import Image
 import io
-from docx import Document
 import fitz  # PyMuPDF
+from docx import Document
 import pandas as pd
-import matplotlib.pyplot as plt
-import seaborn as sns
-import uuid
-from transformers import MarianMTModel, MarianTokenizer
+import logging
+from typing import Optional
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-app = FastAPI()
-
-# Serve static files (HTML, CSS, JS)
+app = FastAPI(title="AI Web Services")
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
-# Initialize models
+# Initialize models (Spaces will cache these)
 try:
-    # Image captioning and VQA
+    logger.info("Loading AI models...")
     image_pipeline = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
-    
-    # Text processing
     text_pipeline = pipeline("text2text-generation", model="t5-small")
-    
-    # Translation models dictionary
-    translation_models = {
-        "fr": "Helsinki-NLP/opus-mt-en-fr",
-        "es": "Helsinki-NLP/opus-mt-en-es",
-        "de": "Helsinki-NLP/opus-mt-en-de"
-    }
-    
-    logger.info("All models loaded successfully")
+    logger.info("Models loaded successfully")
 except Exception as e:
-    logger.error(f"Model loading failed: {str(e)}")
-    raise RuntimeError(f"Failed to initialize models: {str(e)}")
-
-@app.get("/")
-def read_root():
-    return RedirectResponse(url="/static/index.html")
+    logger.error(f"Model loading failed: {e}")
+    raise RuntimeError("Failed to initialize AI models")
 
-@app.get("/health")
-def health_check():
-    return {"status": "healthy", "models_loaded": True}
+@app.get("/", response_class=HTMLResponse)
+async def home():
+    """Serve the frontend interface"""
+    try:
+        with open("static/index.html") as f:
+            return f.read()
+    except Exception as e:
+        logger.error(f"Failed to load frontend: {e}")
+        raise HTTPException(500, "Frontend loading failed")
 
-@app.post("/summarize")
-async def summarize_text(
-    file: UploadFile = File(None),
-    text: str = Form(None)
+@app.post("/api/summarize")
+async def summarize(
+    file: Optional[UploadFile] = File(None),
+    text: Optional[str] = Form(None)
 ):
+    """
+    Summarize text or document
+    Accepts: PDF, DOCX or raw text
+    Returns: {'summary': str}
+    """
     try:
         if file:
-            text = await extract_text_from_file(file)
-        elif not text:
-            raise HTTPException(status_code=400, detail="No input provided")
+            text = await extract_text(file)
+        if not text:
+            raise HTTPException(400, "No text provided")
         
-        summary = text_pipeline(f"summarize: {text}", max_length=100)
-        return {"summary": summary[0]['generated_text']}
+        result = text_pipeline(f"summarize: {text}", max_length=150)
+        return JSONResponse({"summary": result[0]['generated_text']})
+    except HTTPException:
+        raise
     except Exception as e:
-        logger.error(f"Summarization error: {str(e)}")
-        raise HTTPException(status_code=500, detail=str(e))
+        logger.error(f"Summarization error: {e}")
+        raise HTTPException(500, "Summarization failed")
 
-@app.post("/caption")
+@app.post("/api/caption")
 async def caption_image(file: UploadFile = File(...)):
+    """
+    Generate caption for image
+    Accepts: JPEG, PNG
+    Returns: {'caption': str}
+    """
     try:
         image = Image.open(io.BytesIO(await file.read()))
-        caption = image_pipeline(image)
-        return {"caption": caption[0]['generated_text']}
+        result = image_pipeline(image)
+        return JSONResponse({"caption": result[0]['generated_text']})
     except Exception as e:
-        logger.error(f"Captioning error: {str(e)}")
-        raise HTTPException(status_code=500, detail=str(e))
+        logger.error(f"Captioning error: {e}")
+        raise HTTPException(500, "Image captioning failed")
 
-@app.post("/answer")
+@app.post("/api/answer")
 async def answer_question(
-    file: UploadFile = File(None),
-    text: str = Form(None),
+    file: Optional[UploadFile] = File(None),
+    text: Optional[str] = Form(None),
     question: str = Form(...)
 ):
+    """
+    Answer questions about text/document
+    Accepts: PDF, DOCX or raw text + question
+    Returns: {'answer': str}
+    """
     try:
         if file:
-            text = await extract_text_from_file(file)
-        elif not text:
-            raise HTTPException(status_code=400, detail="No context provided")
+            text = await extract_text(file)
+        if not text:
+            raise HTTPException(400, "No text provided")
         
-        answer = text_pipeline(f"question: {question} context: {text}")
-        return {"answer": answer[0]['generated_text']}
-    except Exception as e:
-        logger.error(f"QA error: {str(e)}")
-        raise HTTPException(status_code=500, detail=str(e))
-
-@app.post("/vqa")
-async def visual_question_answering(
-    file: UploadFile = File(...),
-    question: str = Form(...)
-):
-    try:
-        image = Image.open(io.BytesIO(await file.read()))
-        answer = image_pipeline(image, question=question)
-        return {"answer": answer[0]['generated_text']}
+        result = text_pipeline(f"question: {question} context: {text}")
+        return JSONResponse({"answer": result[0]['generated_text']})
+    except HTTPException:
+        raise
     except Exception as e:
-        logger.error(f"VQA error: {str(e)}")
-        raise HTTPException(status_code=500, detail=str(e))
+        logger.error(f"QA error: {e}")
+        raise HTTPException(500, "Question answering failed")
 
-@app.post("/visualize")
-async def visualize_data(
+@app.post("/api/visualize")
+async def generate_visualization(
     file: UploadFile = File(...),
-    request: str = Form(...)
+    chart_type: str = Form("bar")
 ):
+    """
+    Generate visualization code for Excel data
+    Accepts: XLSX, CSV
+    Returns: {'code': str, 'columns': list}
+    """
     try:
         df = pd.read_excel(io.BytesIO(await file.read()))
         
-        if "bar" in request.lower():
+        if chart_type.lower() == "bar":
             code = f"""import matplotlib.pyplot as plt
 plt.bar(df['{df.columns[0]}'], df['{df.columns[1]}'])
+plt.title('Bar Chart')
 plt.show()"""
         else:
             code = f"""import seaborn as sns
 sns.pairplot(df)
+plt.title('Data Distribution')
 plt.show()"""
             
-        return {"code": code}
+        return JSONResponse({
+            "code": code,
+            "columns": list(df.columns)
+        })
     except Exception as e:
-        logger.error(f"Visualization error: {str(e)}")
-        raise HTTPException(status_code=500, detail=str(e))
+        logger.error(f"Visualization error: {e}")
+        raise HTTPException(500, "Visualization code generation failed")
 
-@app.post("/translate")
-async def translate_document(
-    file: UploadFile = File(...),
-    target_language: str = Form(...)
-):
-    try:
-        text = await extract_text_from_file(file)
-        
-        if target_language not in translation_models:
-            raise HTTPException(status_code=400, detail="Unsupported language")
-            
-        tokenizer = MarianTokenizer.from_pretrained(translation_models[target_language])
-        model = MarianMTModel.from_pretrained(translation_models[target_language])
-        
-        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"Translation error: {str(e)}")
-        raise HTTPException(status_code=500, detail=str(e))
-
-async def extract_text_from_file(file: UploadFile):
+async def extract_text(file: UploadFile) -> str:
+    """Extract text from PDF or DOCX files"""
     try:
         content = await file.read()
         
         if file.filename.endswith(".pdf"):
-            doc = fitz.open(stream=content, filetype="pdf")
-            return " ".join([page.get_text() for page in doc])
+            with fitz.open(stream=content, filetype="pdf") as doc:
+                return " ".join(page.get_text() for page in doc)
         elif file.filename.endswith(".docx"):
             doc = Document(io.BytesIO(content))
-            return "\n".join([para.text for para in doc.paragraphs])
+            return "\n".join(p.text for p in doc.paragraphs)
         else:
             raise ValueError("Unsupported file format")
     except Exception as e:
-        logger.error(f"File extraction error: {str(e)}")
-        raise HTTPException(status_code=400, detail=str(e))
+        logger.error(f"Text extraction failed: {e}")
+        raise HTTPException(400, f"Could not extract text: {e}")
 
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+# Health check endpoint
+@app.get("/health")
+async def health_check():
+    return JSONResponse({"status": "healthy", "models": "loaded"})