Update app.py

app.py

@@ -1,4 +1,4 @@
-import os, re
+import os, re, difflib
 from typing import List
 import torch
 import gradio as gr
@@ -8,19 +8,16 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 MODEL_ID = os.getenv("MODEL_ID", "Eemansleepdeprived/Humaneyes")
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Lazy-loaded globals
+# Lazy-loaded
 _tokenizer = None
 _model = None
 
 def load_model():
-    """Load the model on first use (speeds up Space startup)."""
     global _tokenizer, _model
     if _tokenizer is None or _model is None:
         _tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
         _model = AutoModelForSeq2SeqLM.from_pretrained(
-            MODEL_ID,
-            low_cpu_mem_usage=True,
-            torch_dtype=torch.float32,   # CPU safe
+            MODEL_ID, low_cpu_mem_usage=True, torch_dtype=torch.float32
         ).to(device).eval()
     return _tokenizer, _model
 
@@ -45,29 +42,10 @@ def protect(text: str):
 
 def restore(text: str, protected: List[str]):
     def unwrap(m):
-        idx = int(m.group(1))
-        return protected[idx]
+        idx = int(m.group(1)); return protected[idx]
     text = re.sub(rf"{SENTINEL_OPEN}(\d+){SENTINEL_CLOSE}", unwrap, text)
     return text.replace(SENTINEL_OPEN, "").replace(SENTINEL_CLOSE, "")
 
-SYSTEM_RULES = (
-    "Rewrite the text to sound natural, nuanced, and human while preserving meaning.\n"
-    "Rules:\n"
-    "1) Do not change anything between §§KEEP_OPEN§§<id>§§KEEP_CLOSE§§.\n"
-    "2) Keep citations, links, numbers, and code exactly the same.\n"
-    "3) Keep facts the same. Improve clarity, flow, and rhythm. Vary sentence length.\n"
-    "4) No em dashes; use simple punctuation.\n"
-    "5) Keep the requested tone and region.\n"
-)
-
-def build_input(text: str, tone: str, region: str, level: str, intensity: int) -> str:
-    return (
-        f"{SYSTEM_RULES}\n"
-        f"Tone: {tone}. Region: {region} English. Reading level: {level}. "
-        f"Humanization intensity: {intensity} (10 strongest).\n\n"
-        f"INPUT:\n{text}\n\nOUTPUT:\n"
-    )
-
 def chunk_text(s: str, max_chars: int = 1100):
     parts, buf, cur = [], [], 0
     for block in re.split(r"(\n{2,})", s):
@@ -79,33 +57,65 @@ def chunk_text(s: str, max_chars: int = 1100):
     return parts
 
 @torch.inference_mode()
-def generate_one(prompt: str, max_new: int = 256) -> str:
-    tokenizer, model = load_model()
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True).to(device)
-    out_ids = model.generate(
-        **inputs,
-        do_sample=True,
-        temperature=0.4,
-        top_p=0.9,
-        num_beams=1,
-        max_new_tokens=max_new,
-    )
-    return tokenizer.decode(out_ids[0], skip_special_tokens=True)
-
-def humanize_core(text: str, tone: str, region: str, reading_level: str, intensity: int):
+def generate_raw(prompt: str, beams: int = 5, temp: float = None, top_p: float = None, max_new: int = 256) -> str:
+    tok, mdl = load_model()
+    enc = tok(prompt, return_tensors="pt", truncation=True).to(device)
+    gen_kwargs = dict(max_new_tokens=max_new, no_repeat_ngram_size=3)
+    if temp is None:
+        # deterministic beam search (good first try)
+        out = mdl.generate(**enc, num_beams=beams, do_sample=False, **gen_kwargs)
+    else:
+        # stronger change fallback
+        out = mdl.generate(**enc, do_sample=True, temperature=temp, top_p=top_p or 0.9, num_beams=1, **gen_kwargs)
+    return tok.decode(out[0], skip_special_tokens=True)
+
+def difference_ratio(a: str, b: str) -> float:
+    return difflib.SequenceMatcher(None, a, b).ratio()
+
+def humanize_core(text: str, tone: str, region: str, level: str, intensity: int):
+    # 1) Protect spans
     protected_text, bag = protect(text)
-    prompt = build_input(protected_text, tone, region, reading_level, intensity)
-    chunks = chunk_text(prompt, max_chars=1100)
-    outs = [generate_one(ch, max_new=320) for ch in chunks]
-    draft = "".join(outs).replace("—", "-").strip()
+
+    # 2) Create short, model-friendly prompts
+    # Many Pegasus/T5 paraphrasers respond to one of these:
+    prompts = [
+        f"humanize: {protected_text}",
+        f"paraphrase: {protected_text}",
+        protected_text,
+    ]
+
+    # 3) Try deterministic first
+    draft = None
+    for p in prompts:
+        out = generate_raw(p, beams=5, max_new=320)
+        if out.strip():
+            draft = out.strip()
+            if difference_ratio(protected_text, draft) < 0.98:  # changed enough
+                break
+
+    # 4) If barely changed, try a stronger pass (sampling)
+    if draft is None or difference_ratio(protected_text, draft) >= 0.98:
+        for p in prompts:
+            out = generate_raw(p, temp=0.8, top_p=0.92, max_new=320)
+            if out.strip():
+                draft = out.strip()
+                break
+
+    if not draft:
+        draft = protected_text  # absolute fallback
+
+    draft = draft.replace("—", "-")  # enforce simple punctuation
     final_text = restore(draft, bag)
+
+    # paranoia: ensure protected spans survived
     for i, span in enumerate(bag):
         marker = f"{SENTINEL_OPEN}{i}{SENTINEL_CLOSE}"
         if marker in protected_text and span not in final_text:
             final_text = final_text.replace(marker, span)
+
     return final_text
 
-# ---------- Gradio UI (also exposes REST at /api/predict/) ----------
+# ---------- Gradio UI + REST (/api/predict/) ----------
 def ui_humanize(text, tone, region, reading_level, intensity):
     return humanize_core(text, tone, region, reading_level, int(intensity))
 
@@ -120,9 +130,8 @@ demo = gr.Interface(
     ],
     outputs=gr.Textbox(label="Humanized"),
     title="NoteCraft Humanizer (Humaneyes)",
-    description="Model: Eemansleepdeprived/Humaneyes (Pegasus). REST: POST /api/predict/ with { data: [text,tone,region,level,intensity] }",
+    description="REST: POST /api/predict/ with { data: [text,tone,region,level,intensity] }",
 ).queue()
 
 if __name__ == "__main__":
-    # No heavy work here—model loads on first call
     demo.launch()