Update app.py

app.py

@@ -1,34 +1,70 @@
-import os
+import torch
+import pdfplumber
+import pytesseract
+from PIL import Image
+from docx import Document
+from pptx import Presentation
+from transformers import T5Tokenizer, T5ForConditionalGeneration
 from flask import Flask, request, jsonify
-from transformers import pipeline
 
-# Ensure HF doesn't request a token
-os.environ["HF_HOME"] = "/app/cache"
-os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"
-os.environ["HF_HUB_DISABLE_TELEMETRY"] = "1"
-os.environ["HF_HUB_OFFLINE"] = "0"
+# Optimize for CPU
+torch.set_num_threads(4)  # Adjust based on CPU cores
+device = torch.device("cpu")
 
-# Load model
-summarizer = pipeline("summarization", model="t5-base")
+# Load T5-Base model
+model_name = "t5-base"
+tokenizer = T5Tokenizer.from_pretrained(model_name)
+model = T5ForConditionalGeneration.from_pretrained(model_name).to(device)
 
+# Flask App
 app = Flask(__name__)
 
-@app.route("/")
-def home():
-    return "Summarization API is running!"
+# Function to extract text from files
+def extract_text(file):
+    filename = file.filename.lower()
+    
+    if filename.endswith(".pdf"):
+        with pdfplumber.open(file) as pdf:
+            return " ".join([page.extract_text() for page in pdf.pages if page.extract_text()])
+
+    elif filename.endswith(".docx"):
+        doc = Document(file)
+        return " ".join([para.text for para in doc.paragraphs])
+
+    elif filename.endswith(".pptx"):
+        prs = Presentation(file)
+        return " ".join([shape.text for slide in prs.slides for shape in slide.shapes if hasattr(shape, "text")])
+
+    elif filename.endswith((".png", ".jpg", ".jpeg")):
+        image = Image.open(file)
+        return pytesseract.image_to_string(image)
+
+    return None
 
 @app.route("/summarize", methods=["POST"])
-def summarize_text():
-    data = request.get_json()
-    text = data.get("text", "")
-    max_length = data.get("max_length", 50)
-    min_length = data.get("min_length", 10)
+def summarize():
+    file = request.files.get("file")
 
+    if not file:
+        return jsonify({"error": "No file uploaded"}), 400
+
+    text = extract_text(file)
     if not text:
-        return jsonify({"error": "No text provided"}), 400
+        return jsonify({"error": "No text found in file"}), 400
+
+    # Format text for T5
+    input_text = "summarize: " + text.strip()
+
+    # Tokenize input
+    inputs = tokenizer.encode(input_text, return_tensors="pt", truncation=True, max_length=512).to(device)
 
-    summary = summarizer(text, max_length=max_length, min_length=min_length, do_sample=False)
-    return jsonify(summary)
+    # Generate summary
+    with torch.no_grad():
+        summary_ids = model.generate(inputs, max_length=150, min_length=50, length_penalty=2.0, num_beams=4)
+    
+    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+    
+    return jsonify({"summary": summary})
 
 if __name__ == "__main__":
     print("🚀 API is running on port 7860")