Update main.py

main.py

@@ -1,4 +1,4 @@
-"""from fastapi import FastAPI, UploadFile, File, Form, Request
+from fastapi import FastAPI, UploadFile, File, Form, Request
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import HTMLResponse, JSONResponse, FileResponse
 from fastapi.staticfiles import StaticFiles
@@ -72,301 +72,4 @@ async def predict(
         else:
             return JSONResponse({"error": "Invalid option"}, status_code=400)
     except Exception as e:
-        return JSONResponse({"error": f"Prediction failed: {str(e)}"}, status_code=500) """
-from fastapi import FastAPI, UploadFile, File, Form, Request, HTTPException
-from fastapi.responses import HTMLResponse, JSONResponse, FileResponse
-from fastapi.staticfiles import StaticFiles
-from fastapi.templating import Jinja2Templates
-from fastapi.middleware.cors import CORSMiddleware
-from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM, AutoProcessor, AutoModelForCausalLM
-from PIL import Image
-import torch
-import fitz  # PyMuPDF
-import docx
-import pptx
-import openpyxl
-import re
-import nltk
-from nltk.tokenize import sent_tokenize
-from gtts import gTTS
-from fpdf import FPDF
-import tempfile
-import os
-import shutil
-import datetime
-import hashlib
-import easyocr
-from typing import Optional
-
-# Initialize app
-app = FastAPI()
-
-# CORS Configuration
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-# Static assets
-app.mount("/static", StaticFiles(directory="static"), name="static")
-app.mount("/resources", StaticFiles(directory="resources"), name="resources")
-
-# Templates
-templates = Jinja2Templates(directory="templates")
-
-# Initialize models
-nltk.download('punkt', quiet=True)
-
-# Document processing models
-try:
-    tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large-cnn")
-    model = AutoModelForSeq2SeqLM.from_pretrained("facebook/bart-large-cnn")
-    model.eval()
-    summarizer = pipeline("summarization", model=model, tokenizer=tokenizer, device=-1)
-    reader = easyocr.Reader(['en'], gpu=torch.cuda.is_available())
-except Exception as e:
-    print(f"Error loading summarization models: {e}")
-    summarizer = None
-
-# Image captioning models
-try:
-    processor = AutoProcessor.from_pretrained("microsoft/git-large-coco")
-    git_model = AutoModelForCausalLM.from_pretrained("microsoft/git-large-coco")
-    git_model.eval()
-    USE_GIT = True
-except Exception as e:
-    print(f"Error loading GIT model, falling back to ViT: {e}")
-    captioner = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")
-    USE_GIT = False
-
-# Helper functions
-def clean_text(text: str) -> str:
-    text = re.sub(r'\s+', ' ', text)
-    text = re.sub(r'\u2022\s*|\d\.\s+', '', text)
-    text = re.sub(r'\[.*?\]|\(.*?\)', '', text)
-    text = re.sub(r'\bPage\s*\d+\b', '', text, flags=re.IGNORECASE)
-    return text.strip()
-
-def extract_text(file_path: str, file_extension: str):
-    try:
-        if file_extension == "pdf":
-            with fitz.open(file_path) as doc:
-                text = "\n".join(page.get_text("text") for page in doc)
-                if len(text.strip()) < 50:
-                    images = [page.get_pixmap() for page in doc]
-                    temp_img = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
-                    images[0].save(temp_img.name)
-                    ocr_result = reader.readtext(temp_img.name, detail=0)
-                    os.unlink(temp_img.name)
-                    text = "\n".join(ocr_result) if ocr_result else text
-                return clean_text(text), ""
-
-        elif file_extension == "docx":
-            doc = docx.Document(file_path)
-            return clean_text("\n".join(p.text for p in doc.paragraphs)), ""
-
-        elif file_extension == "pptx":
-            prs = pptx.Presentation(file_path)
-            text = [shape.text for slide in prs.slides for shape in slide.shapes if hasattr(shape, "text")]
-            return clean_text("\n".join(text)), ""
-
-        elif file_extension == "xlsx":
-            wb = openpyxl.load_workbook(file_path, read_only=True)
-            text = [" ".join(str(cell) for cell in row if cell) for sheet in wb.sheetnames for row in wb[sheet].iter_rows(values_only=True)]
-            return clean_text("\n".join(text)), ""
-
-        return "", "Unsupported file format"
-    except Exception as e:
-        return "", f"Error reading {file_extension.upper()} file: {str(e)}"
-
-def chunk_text(text: str, max_tokens: int = 950):
-    try:
-        sentences = sent_tokenize(text)
-    except:
-        words = text.split()
-        sentences = [' '.join(words[i:i+20]) for i in range(0, len(words), 20)]
-
-    chunks = []
-    current_chunk = ""
-    for sentence in sentences:
-        token_length = len(tokenizer.encode(current_chunk + " " + sentence))
-        if token_length <= max_tokens:
-            current_chunk += " " + sentence
-        else:
-            chunks.append(current_chunk.strip())
-            current_chunk = sentence
-
-    if current_chunk:
-        chunks.append(current_chunk.strip())
-
-    return chunks
-
-def generate_summary(text: str, length: str = "medium") -> str:
-    cache_key = hashlib.md5((text + length).encode()).hexdigest()
-    
-    length_params = {
-        "short": {"max_length": 80, "min_length": 30},
-        "medium": {"max_length": 200, "min_length": 80},
-        "long": {"max_length": 300, "min_length": 210}
-    }
-    
-    chunks = chunk_text(text)
-    try:
-        summaries = summarizer(
-            chunks,
-            max_length=length_params[length]["max_length"],
-            min_length=length_params[length]["min_length"],
-            do_sample=False,
-            truncation=True
-        )
-        summary_texts = [s['summary_text'] for s in summaries]
-    except Exception as e:
-        summary_texts = [f"[Error: {str(e)}"]
-
-    final_summary = " ".join(summary_texts)
-    final_summary = ". ".join(s.strip().capitalize() for s in final_summary.split(". ") if s.strip())
-    return final_summary if len(final_summary) > 25 else "Summary too short"
-
-def text_to_speech(text: str):
-    try:
-        tts = gTTS(text)
-        temp_audio = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
-        tts.save(temp_audio.name)
-        return temp_audio.name
-    except Exception as e:
-        print(f"Error in text-to-speech: {e}")
-        return ""
-
-def create_pdf(summary: str, original_filename: str):
-    try:
-        pdf = FPDF()
-        pdf.add_page()
-        pdf.set_font("Arial", 'B', 16)
-        pdf.cell(200, 10, txt="Document Summary", ln=1, align='C')
-        pdf.set_font("Arial", size=12)
-        pdf.cell(200, 10, txt=f"Original file: {original_filename}", ln=1)
-        pdf.cell(200, 10, txt=f"Generated on: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}", ln=1)
-        pdf.ln(10)
-        pdf.multi_cell(0, 10, txt=summary)
-        temp_pdf = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf")
-        pdf.output(temp_pdf.name)
-        return temp_pdf.name
-    except Exception as e:
-        print(f"Error creating PDF: {e}")
-        return ""
-
-def generate_caption(image_path: str) -> str:
-    try:
-        if USE_GIT:
-            image = Image.open(image_path).convert("RGB")
-            inputs = processor(images=image, return_tensors="pt")
-            outputs = git_model.generate(**inputs, max_length=50)
-            caption = processor.batch_decode(outputs, skip_special_tokens=True)[0]
-        else:
-            result = captioner(image_path)
-            caption = result[0]['generated_text']
-        return caption
-    except Exception as e:
-        raise Exception(f"Caption generation failed: {str(e)}")
-
-# API Endpoints
-@app.post("/summarize/")
-async def summarize_document(file: UploadFile = File(...), length: str = Form("medium")):
-    valid_types = [
-        'application/pdf',
-        'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
-        'application/vnd.openxmlformats-officedocument.presentationml.presentation',
-        'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'
-    ]
-    
-    if file.content_type not in valid_types:
-        raise HTTPException(
-            status_code=400,
-            detail="Please upload a valid document (PDF, DOCX, PPTX, or XLSX)"
-        )
-
-    try:
-        # Save temp file
-        with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(file.filename)[1]) as temp:
-            shutil.copyfileobj(file.file, temp)
-            temp_path = temp.name
-
-        # Process file
-        text, error = extract_text(temp_path, os.path.splitext(file.filename)[1][1:].lower())
-        if error:
-            raise HTTPException(status_code=400, detail=error)
-
-        if not text or len(text.split()) < 30:
-            raise HTTPException(status_code=400, detail="Document too short to summarize")
-
-        summary = generate_summary(text, length)
-        audio_path = text_to_speech(summary)
-        pdf_path = create_pdf(summary, file.filename)
-
-        return {
-            "summary": summary,
-            "audio_url": f"/files/{os.path.basename(audio_path)}" if audio_path else None,
-            "pdf_url": f"/files/{os.path.basename(pdf_path)}" if pdf_path else None
-        }
-
-    except HTTPException:
-        raise
-    except Exception as e:
-        raise HTTPException(
-            status_code=500,
-            detail=f"Summarization failed: {str(e)}"
-        )
-    finally:
-        if 'temp_path' in locals() and os.path.exists(temp_path):
-            os.unlink(temp_path)
-
-@app.post("/imagecaption/")
-async def caption_image(file: UploadFile = File(...)):
-    valid_types = ['image/jpeg', 'image/png', 'image/gif', 'image/webp']
-    if file.content_type not in valid_types:
-        raise HTTPException(
-            status_code=400,
-            detail="Please upload a valid image (JPEG, PNG, GIF, or WEBP)"
-        )
-
-    try:
-        # Save temp file
-        with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(file.filename)[1]) as temp:
-            shutil.copyfileobj(file.file, temp)
-            temp_path = temp.name
-
-        # Generate caption
-        caption = generate_caption(temp_path)
-        
-        # Generate audio
-        audio_path = text_to_speech(caption)
-
-        return {
-            "answer": caption,
-            "audio": f"/files/{os.path.basename(audio_path)}" if audio_path else None
-        }
-
-    except HTTPException:
-        raise
-    except Exception as e:
-        raise HTTPException(
-            status_code=500,
-            detail=str(e)
-        )
-    finally:
-        if 'temp_path' in locals() and os.path.exists(temp_path):
-            os.unlink(temp_path)
-
-@app.get("/files/{filename}")
-async def serve_file(filename: str):
-    file_path = os.path.join(tempfile.gettempdir(), filename)
-    if os.path.exists(file_path):
-        return FileResponse(file_path)
-    raise HTTPException(status_code=404, detail="File not found")
-
-@app.get("/", response_class=HTMLResponse)
-async def serve_home(request: Request):
-    return templates.TemplateResponse("HomeS.html", {"request": request})
+        return JSONResponse({"error": f"Prediction failed: {str(e)}"}, status_code=500)