Update main.py

main.py

@@ -3,16 +3,28 @@ os.environ['HF_HOME'] = '/tmp'
 
 from flask import Flask, request, jsonify, render_template
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
-import math, textwrap
+import json, re, time
 
 app = Flask(__name__)
 
-MODEL_NAME = "sshleifer/distilbart-cnn-12-6"
+# -------------------------
+# Models (CPU as requested)
+# -------------------------
+# Primary summarizer: higher-quality model
+MODEL_NAME = "facebook/bart-large-cnn"
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
-summarizer = pipeline("summarization", model=model, tokenizer=tokenizer, device=-1)  # set device appropriately
+summarizer = pipeline("summarization", model=model, tokenizer=tokenizer, device=-1)  # CPU
 
-# Simple mapping of presets to generation lengths
+# Small instruction model to choose length/tone when "auto" is requested
+PARAM_MODEL_NAME = "google/flan-t5-small"
+param_tokenizer = AutoTokenizer.from_pretrained(PARAM_MODEL_NAME)
+param_model = AutoModelForSeq2SeqLM.from_pretrained(PARAM_MODEL_NAME)
+param_generator = pipeline("text2text-generation", model=param_model, tokenizer=param_tokenizer, device=-1)  # CPU
+
+# -------------------------
+# Presets & helpers
+# -------------------------
 LENGTH_PRESETS = {
     "short": {"min_length": 20, "max_length": 60},
     "medium": {"min_length": 60, "max_length": 130},
@@ -47,31 +59,127 @@ def apply_tone_instruction(text, tone):
         instr = "Summarize:"
     return f"{instr}\n\n{text}"
 
-# NEW: Route to show summarizer.html (fixes 404)
+# small regex int extractor
+def _first_int_from_text(s, fallback=None):
+    m = re.search(r"\d{1,5}", s)
+    return int(m.group()) if m else fallback
+
+def generate_summarization_config(text):
+    """
+    Uses the small instruction model to recommend:
+      - length: short|medium|long
+      - min_words, max_words (integers)
+      - tone: neutral|formal|casual|bullet
+    Returns a normalized dict with keys: length, min_length, max_length, tone
+    Falls back to heuristics on failure.
+    """
+    prompt = (
+        "You are a helpful assistant that recommends summarization settings.\n"
+        "Given the following source text, pick a summary LENGTH category (short/medium/long), "
+        "an estimated MIN and MAX length in words for the summary, and a TONE (neutral/formal/casual/bullet).\n"
+        "Respond in a single line in this exact JSON format (no extra commentary):\n"
+        '{"length":"SHORT_OR_MEDIUM_OR_LONG","min_words":MIN,"max_words":MAX,"tone":"NEUTRAL|FORMAL|CASUAL|BULLET"}\n\n'
+        "Text:\n'''"
+        + (text[:6000])
+        + "'''"
+    )
+    try:
+        out = param_generator(prompt, max_length=200, do_sample=False)[0]["generated_text"].strip()
+        # attempt parse
+        try:
+            cfg = json.loads(out)
+        except Exception:
+            jmatch = re.search(r"\{.*\}", out, re.DOTALL)
+            if jmatch:
+                raw = jmatch.group()
+                raw = raw.replace("'", '"')
+                cfg = json.loads(raw)
+            else:
+                raise
+
+        length = cfg.get("length", "").lower()
+        tone = cfg.get("tone", "").lower()
+        min_w = cfg.get("min_words")
+        max_w = cfg.get("max_words")
+
+        # sensible defaults if parse odd
+        if length not in ("short", "medium", "long"):
+            words = len(text.split())
+            length = "short" if words < 150 else ("medium" if words < 800 else "long")
+        if tone not in ("neutral", "formal", "casual", "bullet"):
+            tone = "neutral"
+
+        # fallback numeric extraction
+        if not isinstance(min_w, int):
+            min_w = _first_int_from_text(out, fallback=None)
+        if not isinstance(max_w, int):
+            max_w = _first_int_from_text(out[::-1], fallback=None)
+
+        defaults = {"short": (15, 50), "medium": (50, 130), "long": (130, 300)}
+        dmin, dmax = defaults.get(length, (50, 130))
+        min_len = int(min_w) if isinstance(min_w, int) else dmin
+        max_len = int(max_w) if isinstance(max_w, int) else dmax
+
+        # clamp to sane bounds
+        min_len = max(5, min(min_len, 2000))
+        max_len = max(min_len + 5, min(max_len, 4000))
+
+        return {"length": length, "min_length": min_len, "max_length": max_len, "tone": tone}
+    except Exception:
+        # fallback heuristic
+        words = len(text.split())
+        length = "short" if words < 150 else ("medium" if words < 800 else "long")
+        d = {"short": (15, 50), "medium": (50, 130), "long": (130, 300)}
+        mn, mx = d[length]
+        return {"length": length, "min_length": mn, "max_length": mx, "tone": "neutral"}
+
+# -------------------------
+# Routes
+# -------------------------
 @app.route("/")
 def home():
+    # expects templates/index.html to exist (your frontend)
     return render_template("index.html")
 
 @app.route("/summarize", methods=["POST"])
 def summarize_route():
+    start_time = time.time()
     data = request.get_json(force=True)
-    text = data.get("text", "")[:20000]
-    length = data.get("length", "medium")
-    tone = data.get("tone", "neutral")
+    text = data.get("text", "")[:20000]  # cap input
+    requested_length = (data.get("length") or "medium").lower()
+    requested_tone = (data.get("tone") or "neutral").lower()
 
     if not text or len(text.split()) < 5:
         return jsonify({"error": "Input too short."}), 400
 
-    preset = LENGTH_PRESETS.get(length, LENGTH_PRESETS["medium"])
+    # If user asks AI to choose settings
+    if requested_length in ("auto", "ai") or requested_tone in ("auto", "ai"):
+        cfg = generate_summarization_config(text)
+        length = cfg.get("length", "medium")
+        tone = cfg.get("tone", "neutral")
+        preset_min = cfg.get("min_length")
+        preset_max = cfg.get("max_length")
+        preset = LENGTH_PRESETS.get(length, LENGTH_PRESETS["medium"])
+    else:
+        length = requested_length if requested_length in LENGTH_PRESETS else "medium"
+        tone = requested_tone if requested_tone in ("neutral", "formal", "casual", "bullet") else "neutral"
+        preset = LENGTH_PRESETS.get(length, LENGTH_PRESETS["medium"])
+        preset_min = preset["min_length"]
+        preset_max = preset["max_length"]
+
+    # chunk input for long texts
     chunks = chunk_text_by_chars(text, max_chars=1500, overlap=200)
     summaries = []
 
     for chunk in chunks:
         prompted = apply_tone_instruction(chunk, tone)
+        min_l = int(preset_min) if preset_min is not None else preset["min_length"]
+        max_l = int(preset_max) if preset_max is not None else preset["max_length"]
+
         out = summarizer(
             prompted,
-            min_length=preset["min_length"],
-            max_length=preset["max_length"],
+            min_length=min_l,
+            max_length=max_l,
             truncation=True
         )[0]["summary_text"]
         summaries.append(out.strip())
@@ -92,7 +200,9 @@ def summarize_route():
         lines = [l.strip() for s in final.splitlines() for l in s.split(". ") if l.strip()]
         final = "\n".join(f"- {l.rstrip('.')}" for l in lines[:20])
 
-    return jsonify({"summary": final})
+    elapsed = time.time() - start_time
+    return jsonify({"summary": final, "meta": {"length_choice": length, "tone": tone, "time_seconds": round(elapsed, 2)}})
 
 if __name__ == "__main__":
-    app.run(debug=True, port=7860)
+    # keep debug off in production; using CPU as requested
+    app.run(host="0.0.0.0", port=7860, debug=True)