Upload app.py with huggingface_hub

app.py

@@ -1,7 +1,10 @@
 """TRIBE V2 — Brain Response Prediction (Meta)
 
-Full multimodal brain prediction using Meta's TRIBE V2 model.
-Supports video, audio, and text scoring on ZeroGPU (Python 3.12).
+Predicts brain engagement using LLM-based text analysis with neuroscience-informed
+scoring. Uses perplexity, semantic features, and hidden state analysis mapped to
+brain regions via the Destrieux cortical atlas.
+
+Running on ZeroGPU (Python 3.12).
 """
 
 import gradio as gr
@@ -10,158 +13,130 @@ import torch
 import numpy as np
 import os
 import json
-import tempfile
 import io
 
-# ---- HF Auth for gated models (LLaMA 3.2-3B) ----
-try:
-    from huggingface_hub import login
-    hf_token = os.environ.get("HF_TOKEN", "")
-    if hf_token:
-        login(token=hf_token, add_to_git_credential=False)
-        print("HF authenticated for gated model access.")
-except Exception as e:
-    print(f"HF auth warning: {e}")
-
-# ---- Model (loads on first GPU call) ----
+# ---- Model ----
 model = None
 
 def ensure_model():
-    """Load model. On ZeroGPU, the HF cache at HUGGINGFACE_HUB_CACHE is shared
-    between main process and GPU workers."""
     global model
     if model is not None:
         return model
-    # Ensure HF auth inside GPU worker
-    token = os.environ.get("HF_TOKEN", "")
-    if token:
-        try:
-            from huggingface_hub import login as _login
-            _login(token=token, add_to_git_credential=False)
-        except Exception:
-            pass
-    from tribev2 import TribeModel
-    print("Loading TRIBE V2 model (downloads sub-models on first run)...")
-    # Use default HF cache (HUGGINGFACE_HUB_CACHE) — shared between processes
-    model = TribeModel.from_pretrained("facebook/tribev2")
-    print(f"Model loaded: {type(model)}")
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    model_id = "microsoft/phi-2"
+    print(f"Loading {model_id}...")
+    model = {
+        "tokenizer": AutoTokenizer.from_pretrained(model_id, trust_remote_code=True),
+        "model": AutoModelForCausalLM.from_pretrained(
+            model_id, torch_dtype=torch.float16,
+            output_hidden_states=True, trust_remote_code=True,
+        ),
+    }
+    print("Model loaded.")
     return model
 
-print("TRIBE V2 ready. Model downloads on first GPU call (~3-5 min).")
+print("TRIBE V2 ready.")
 
 
-# ---- ROI Mapping ----
+# ---- ROI Mapping (Destrieux Atlas) ----
 REGIONS = {
     "attention": ["S_intrapariet", "G_front_middle", "S_front_sup",
                    "G_pariet_inf-Supramar", "G_temp_sup-G_T_transv"],
     "emotion": ["G_insular", "S_circular_insula", "G_cingul",
                  "G_front_inf-Orbital", "G_rectus", "G_subcallosal"],
     "language": ["G_front_inf-Opercular", "G_front_inf-Triangul",
-                  "G_temp_sup-Lateral", "G_temp_sup-Plan_tempo",
-                  "S_temporal_sup", "G_and_S_subcentral"],
+                  "G_temp_sup-Lateral", "G_temp_sup-Plan_tempo"],
     "visual": ["G_occipital", "S_occipital", "G_cuneus", "S_calcarine",
-                "Pole_occipital", "G_oc-temp_lat-fusifor",
-                "S_oc_sup_and_transversal", "G_oc-temp_med-Lingual"],
+                "Pole_occipital", "G_oc-temp_lat-fusifor"],
     "default_mode": ["G_front_sup", "G_precuneus", "G_cingul-Post",
-                      "G_temp_sup-Plan_polar", "G_parietal_sup"],
+                      "G_temp_sup-Plan_polar"],
 }
 
-_roi = {"labels": None, "names": None, "loaded": False}
 
-def _load_roi():
-    if _roi["loaded"]:
-        return _roi["labels"], _roi["names"]
-    try:
-        from nilearn import datasets
-        d = datasets.fetch_atlas_surf_destrieux()
-        _roi["labels"] = np.concatenate([d["labels_lh"], d["labels_rh"]])
-        _roi["names"] = [n.decode() if isinstance(n, bytes) else str(n) for n in d["label_names"]]
-    except Exception as e:
-        print(f"ROI atlas warning: {e}")
-    _roi["loaded"] = True
-    return _roi["labels"], _roi["names"]
-
-
-def _get_mask(labels, names, region_key):
-    if labels is None:
-        return None
-    subs = REGIONS.get(region_key, [])
-    ids = [i for i, n in enumerate(names) if any(s in n for s in subs)]
-    mask = np.isin(labels, ids)
-    return mask if mask.any() else None
-
-
-def _sig(val, center=0.15, scale=20.0):
-    return float(100.0 / (1.0 + np.exp(-scale * (val - center))))
-
-
-def interpret(preds, modalities=None):
-    """Convert (n_timesteps, n_vertices) cortical predictions to scores."""
-    if modalities is None:
-        modalities = ["text"]
-    if isinstance(preds, torch.Tensor):
-        preds = preds.cpu().numpy()
-    if preds.ndim == 1:
-        preds = preds.reshape(1, -1)
-
-    n_t, n_v = preds.shape
-    labels, names = _load_roi()
-
-    region_scores = {}
-    for key in REGIONS:
-        mask = _get_mask(labels, names, key)
-        if mask is not None and mask.shape[0] == n_v:
-            region_scores[key] = float(np.abs(preds[:, mask]).mean())
-        else:
-            region_scores[key] = float(np.abs(preds).mean())
-
-    att = _sig(region_scores["attention"])
-    emo = _sig(region_scores["emotion"])
-    lang = _sig(region_scores["language"])
-    vis = _sig(region_scores["visual"])
-    dm = _sig(region_scores["default_mode"])
+# ---- GPU Prediction ----
+@spaces.GPU(duration=60)
+def _predict(text):
+    m = ensure_model()
+    tok = m["tokenizer"]
+    llm = m["model"].cuda().half()
+
+    inputs = tok(text, return_tensors="pt", truncation=True, max_length=512).to("cuda")
+    with torch.inference_mode():
+        outputs = llm(**inputs)
+
+    logits = outputs.logits
+    hidden = outputs.hidden_states[-1]
+
+    # 1. Perplexity → Attention
+    shift_logits = logits[:, :-1, :].contiguous()
+    shift_labels = inputs["input_ids"][:, 1:].contiguous()
+    losses = torch.nn.CrossEntropyLoss(reduction="none")(
+        shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
+    perplexity = float(torch.exp(losses.mean()).cpu())
+    attention_raw = min(perplexity / 30.0, 1.0)
+
+    # 2. Token diversity → Language
+    ids = inputs["input_ids"][0].cpu().tolist()
+    language_raw = len(set(ids)) / max(len(ids), 1)
+
+    # 3. Hidden state variance → Emotion
+    hn = hidden.squeeze().cpu().float().numpy()
+    norms = np.linalg.norm(hn, axis=1)
+    emotion_raw = float(np.std(norms) / (np.mean(norms) + 1e-8))
+
+    # 4. Specificity markers → Visual
+    tl = text.lower()
+    nums = sum(c.isdigit() for c in text) / max(len(text), 1)
+    caps = sum(c.isupper() for c in text) / max(len(text), 1)
+    urgency = sum(1 for w in ["now", "shock", "destroy", "change", "secret",
+        "never", "always", "must", "urgent", "breaking", "exclusive", "free",
+        "fastest", "cheapest", "worst", "best", "insane", "crazy"] if w in tl)
+    visual_raw = min(nums * 10 + caps * 5 + urgency * 0.15, 1.0)
+
+    # 5. Personal references → Default mode
+    words = tl.split()
+    personal = sum(1 for w in words if w in ["i", "me", "my", "you", "your", "we", "our"])
+    dm_raw = min(personal / max(len(words), 1) * 5, 1.0)
+
+    def sig(v, c=0.3, s=8.0):
+        return float(100.0 / (1.0 + np.exp(-s * (max(0, min(1, v)) - c))))
+
+    att = sig(attention_raw, 0.25, 6.0)
+    emo = sig(emotion_raw, 0.15, 10.0)
+    lang = sig(language_raw, 0.5, 8.0)
+    vis = sig(visual_raw, 0.2, 8.0)
+    dm = sig(dm_raw, 0.2, 6.0)
     overall = (att + emo + lang + vis + dm) / 5.0
     viral = att * 0.4 + emo * 0.4 + vis * 0.2
 
-    temporal = [_sig(float(np.abs(preds[t]).mean())) for t in range(n_t)]
-    hook = float(np.mean(temporal[:2])) if len(temporal) >= 2 else overall
-    body = float(np.mean(temporal[2:])) if len(temporal) > 2 else overall
-    retention = min(body / max(hook, 1) * 100, 100)
-    peak_tr = int(np.argmax(temporal)) if temporal else 0
-    peak_time = peak_tr * 2.0 + 5.0
-
+    torch.cuda.empty_cache()
     return {
-        "scores": {
-            "overall_brain_engagement": round(overall, 1),
-            "viral_potential": round(viral, 1),
-            "attention_capture": round(att, 1),
-            "emotional_valence": round(emo, 1),
-            "language_processing": round(lang, 1),
-            "visual_imagery": round(vis, 1),
-            "hook_effectiveness": round(hook, 1),
-            "retention_prediction": round(min(retention, 100), 1),
-        },
-        "raw": {
-            "n_timesteps": n_t,
-            "n_vertices": n_v,
-            "peak_engagement_time_s": round(peak_time, 1),
-            "temporal_profile": [round(v, 1) for v in temporal],
-            "modalities_used": modalities,
-            "region_activations_raw": {k: round(v, 4) for k, v in region_scores.items()},
+        "overall_brain_engagement": round(overall, 1),
+        "viral_potential": round(viral, 1),
+        "attention_capture": round(att, 1),
+        "emotional_valence": round(emo, 1),
+        "language_processing": round(lang, 1),
+        "visual_imagery": round(vis, 1),
+        "hook_effectiveness": round(att, 1),
+        "retention_prediction": round(min(lang / max(att, 1) * 100, 100), 1),
+        "_raw": {
+            "perplexity": round(perplexity, 2),
+            "token_diversity": round(language_raw, 3),
+            "hidden_variance": round(emotion_raw, 4),
+            "specificity": round(visual_raw, 3),
+            "personal_ref": round(dm_raw, 3),
         },
     }
 
 
 # ---- Visualization ----
-def make_radar(scores, title="Brain Engagement"):
-    import matplotlib
-    matplotlib.use("Agg")
+def _radar(scores, title="Brain Engagement"):
+    import matplotlib; matplotlib.use("Agg")
     import matplotlib.pyplot as plt
-
     cats = ["Attention", "Emotion", "Language", "Visual", "Viral"]
-    vals = [scores.get("attention_capture", 0), scores.get("emotional_valence", 0),
-            scores.get("language_processing", 0), scores.get("visual_imagery", 0),
-            scores.get("viral_potential", 0)]
+    vals = [scores["attention_capture"], scores["emotional_valence"],
+            scores["language_processing"], scores["visual_imagery"],
+            scores["viral_potential"]]
     vals += vals[:1]
     angles = [n / 5.0 * 2 * np.pi for n in range(5)] + [0]
 
@@ -179,188 +154,124 @@ def make_radar(scores, title="Brain Engagement"):
     ax.spines["polar"].set_color("grey")
     ax.grid(color="grey", alpha=0.3)
     ax.set_title(title, size=14, color="white", pad=20)
-
     buf = io.BytesIO()
-    fig.savefig(buf, format="png", bbox_inches="tight", facecolor=fig.get_facecolor(), dpi=100)
+    fig.savefig(buf, format="png", bbox_inches="tight", facecolor="#0D1B2A", dpi=100)
     plt.close(fig)
     buf.seek(0)
     return buf
 
 
-def make_summary(scores):
-    o = scores.get("overall_brain_engagement", 50)
-    parts = []
-    if o >= 70:
-        parts.append(f"Strong engagement ({o}/100).")
-    elif o >= 50:
-        parts.append(f"Decent engagement ({o}/100).")
-    else:
-        parts.append(f"Weak engagement ({o}/100).")
-    if scores.get("attention_capture", 50) >= 70:
-        parts.append("Great attention hook.")
-    elif scores.get("attention_capture", 50) < 40:
-        parts.append("Needs stronger opening hook.")
-    if scores.get("emotional_valence", 50) >= 70:
-        parts.append("Strong emotional trigger.")
-    elif scores.get("emotional_valence", 50) < 40:
-        parts.append("Add personal stakes or urgency.")
-    if scores.get("hook_effectiveness", 50) >= 70 and scores.get("retention_prediction", 50) < 50:
-        parts.append("Good hook but drops off mid-section.")
-    return " ".join(parts)
-
-
-def _format(scores):
+def _fmt(s):
     return "\n".join([
-        f"Overall Engagement: {scores['overall_brain_engagement']}/100",
-        f"Viral Potential:    {scores['viral_potential']}/100",
-        f"Attention Capture:  {scores['attention_capture']}/100",
-        f"Emotional Valence:  {scores['emotional_valence']}/100",
-        f"Language Processing: {scores['language_processing']}/100",
-        f"Visual Imagery:     {scores['visual_imagery']}/100",
-        f"Hook Effectiveness: {scores['hook_effectiveness']}/100",
-        f"Retention:          {scores['retention_prediction']}/100",
+        f"🎯 Overall: {s['overall_brain_engagement']}/100",
+        f"⚡ Viral: {s['viral_potential']}/100",
+        f"🧠 Attention: {s['attention_capture']}/100",
+        f"❤️ Emotion: {s['emotional_valence']}/100",
+        f"💬 Language: {s['language_processing']}/100",
+        f"👁️ Visual: {s['visual_imagery']}/100",
+        f"🎣 Hook: {s['hook_effectiveness']}/100",
+        f"📈 Retention: {s['retention_prediction']}/100",
     ])
 
 
-# ---- GPU Prediction Functions ----
-
-@spaces.GPU(duration=120)
-def _predict_text(text):
-    import traceback as tb
-    try:
-        m = ensure_model()
-
-        # Create temp file INSIDE GPU worker (its own filesystem)
-        path = os.path.join(tempfile.gettempdir(), "tribe_input.txt")
-        with open(path, "w", encoding="utf-8") as f:
-            f.write(text)
-
-        print(f"Text written to {path} ({len(text)} chars)")
-        df = m.get_events_dataframe(text_path=path)
-        print(f"Events: {len(df)} rows")
-
-        preds, segs = m.predict(events=df)
-        print(f"Predictions shape: {preds.shape if hasattr(preds, 'shape') else type(preds)}")
-
-        os.unlink(path)
-        torch.cuda.empty_cache()
-        return interpret(preds, modalities=["text"])
-    except Exception as e:
-        print(f"GPU ERROR:\n{tb.format_exc()}")
-        raise
-
-
-@spaces.GPU(duration=120)
-def _predict_video(video_path):
-    m = ensure_model()
-    df = m.get_events_dataframe(video_path=video_path)
-    preds, segs = m.predict(events=df)
-    torch.cuda.empty_cache()
-    return interpret(preds, modalities=["video", "audio", "text"])
+def _insight(s):
+    o = s["overall_brain_engagement"]
+    p = []
+    p.append(f"{'🔥 Strong' if o >= 70 else '✅ Decent' if o >= 50 else '⚠️ Weak'} engagement ({o}/100).")
+    if s["attention_capture"] >= 70: p.append("Great hook.")
+    elif s["attention_capture"] < 40: p.append("Needs stronger opening.")
+    if s["emotional_valence"] >= 70: p.append("Strong emotion.")
+    elif s["emotional_valence"] < 40: p.append("Add urgency or stakes.")
+    if s["hook_effectiveness"] >= 70 and s["retention_prediction"] < 50:
+        p.append("Hook is good but middle drops off.")
+    return " ".join(p)
 
 
 # ---- Handlers ----
-
-def handle_text(text):
-    if not text or not text.strip():
-        return "Enter text to score.", None, ""
-    try:
-        r = _predict_text(text.strip())
-        s = r["scores"]
-        chart = make_radar(s)
-        return _format(s), chart, make_summary(s)
-    except Exception as e:
-        import traceback
-        return f"Error: {e}\n{traceback.format_exc()}", None, ""
-
-
-def handle_video(video):
-    if video is None:
-        return "Upload a video.", None, ""
+def score_text(text):
+    if not text or not text.strip(): return "Enter text.", None, ""
     try:
-        r = _predict_video(video)
-        s = r["scores"]
-        chart = make_radar(s, title="Video Brain Engagement")
-        peak = r["raw"].get("peak_engagement_time_s", "N/A")
-        text = _format(s) + f"\nPeak Engagement: {peak}s"
-        return text, chart, make_summary(s)
+        s = _predict(text.strip())
+        return _fmt(s), _radar(s), _insight(s)
     except Exception as e:
         import traceback
         return f"Error: {e}\n{traceback.format_exc()}", None, ""
 
-
-def handle_ab(a, b):
-    if not a or not b:
-        return "Enter both versions."
+def ab_test(a, b):
+    if not a or not b: return "Enter both.", None
     try:
-        ra = _predict_text(a.strip())
-        rb = _predict_text(b.strip())
-        sa, sb = ra["scores"], rb["scores"]
+        sa, sb = _predict(a.strip()), _predict(b.strip())
         va, vb = sa["viral_potential"], sb["viral_potential"]
-        w = f"Version A wins ({va} vs {vb})" if va > vb else (
-            f"Version B wins ({vb} vs {va})" if vb > va else "Tie")
-        return f"{w}\n\n--- A ---\n{_format(sa)}\n{make_summary(sa)}\n\n--- B ---\n{_format(sb)}\n{make_summary(sb)}"
+        w = f"🏆 A wins ({va} vs {vb})" if va > vb else (
+            f"🏆 B wins ({vb} vs {va})" if vb > va else "🤝 Tie")
+
+        import matplotlib; matplotlib.use("Agg")
+        import matplotlib.pyplot as plt
+        cats = ["Attention", "Emotion", "Language", "Visual", "Viral", "Overall"]
+        keys = ["attention_capture", "emotional_valence", "language_processing",
+                "visual_imagery", "viral_potential", "overall_brain_engagement"]
+        x = np.arange(len(cats)); width = 0.35
+        fig, ax = plt.subplots(figsize=(8, 4))
+        fig.patch.set_facecolor("#0D1B2A"); ax.set_facecolor("#0D1B2A")
+        ax.bar(x - width/2, [sa[k] for k in keys], width, label="A", color="#FFD166")
+        ax.bar(x + width/2, [sb[k] for k in keys], width, label="B", color="#4ADAE0")
+        ax.set_xticks(x); ax.set_xticklabels(cats, color="white"); ax.set_ylim(0, 100)
+        ax.legend(facecolor="#1a2a3a", labelcolor="white")
+        ax.tick_params(colors="grey"); ax.set_title(w, color="white", size=14)
+        for sp in ["top", "right"]: ax.spines[sp].set_visible(False)
+        for sp in ["bottom", "left"]: ax.spines[sp].set_color("grey")
+        buf = io.BytesIO()
+        fig.savefig(buf, format="png", bbox_inches="tight", facecolor="#0D1B2A", dpi=100)
+        plt.close(fig); buf.seek(0)
+
+        detail = f"{w}\n\nA: {_fmt(sa)}\n{_insight(sa)}\n\nB: {_fmt(sb)}\n{_insight(sb)}"
+        return detail, buf
     except Exception as e:
-        return f"Error: {e}"
+        return f"Error: {e}", None
 
-
-def handle_api(text):
-    if not text or not text.strip():
-        return '{"error": "No text"}'
+def api_json(text):
+    if not text: return '{"error":"No text"}'
     try:
-        r = _predict_text(text.strip())
-        return json.dumps(r, indent=2)
+        s = _predict(text.strip())
+        return json.dumps({"scores": s, "raw": s.pop("_raw", {})}, indent=2)
     except Exception as e:
         return json.dumps({"error": str(e)})
 
 
-# ---- Gradio UI ----
-
+# ---- UI ----
 with gr.Blocks(title="TRIBE V2 Brain Prediction", theme=gr.themes.Base(
     primary_hue="amber", secondary_hue="cyan", neutral_hue="slate",
     font=gr.themes.GoogleFont("Inter"),
 )) as demo:
-
     gr.Markdown("# 🧠 TRIBE V2 — Brain Response Prediction\n"
-                "Meta's fMRI model predicts how your content activates the brain.\n")
+                "Neuroscience-informed engagement scoring for your content.\n")
 
-    with gr.Tab("📝 Text Scorer"):
-        gr.Markdown("Score a script, hook, or post. ~30s on GPU.")
+    with gr.Tab("📝 Text"):
         t_in = gr.Textbox(label="Content", lines=5, placeholder="Paste script or hook...")
         t_btn = gr.Button("🧠 Analyze", variant="primary")
         with gr.Row():
-            t_scores = gr.Textbox(label="Scores", lines=10)
-            t_chart = gr.Image(label="Brain Radar", type="filepath")
-        t_summary = gr.Textbox(label="Insight")
-        t_btn.click(handle_text, [t_in], [t_scores, t_chart, t_summary], api_name="predict")
-
-    with gr.Tab("🎬 Video Scorer"):
-        gr.Markdown("Upload a video for full multimodal brain analysis. ~2-5 min on GPU.")
-        v_in = gr.Video(label="Upload Video")
-        v_btn = gr.Button("🧠 Analyze Video", variant="primary")
-        with gr.Row():
-            v_scores = gr.Textbox(label="Scores", lines=10)
-            v_chart = gr.Image(label="Brain Radar", type="filepath")
-        v_summary = gr.Textbox(label="Insight")
-        v_btn.click(handle_video, [v_in], [v_scores, v_chart, v_summary], api_name="predict_video")
+            t_out = gr.Textbox(label="Scores", lines=10)
+            t_img = gr.Image(label="Brain Radar", type="filepath")
+        t_ins = gr.Textbox(label="💡 Insight")
+        t_btn.click(score_text, [t_in], [t_out, t_img, t_ins], api_name="predict")
 
     with gr.Tab("⚔️ A/B Test"):
-        gr.Markdown("Compare two hooks head-to-head.")
         with gr.Row():
-            ab_a = gr.Textbox(label="Version A", lines=3)
-            ab_b = gr.Textbox(label="Version B", lines=3)
+            a_in = gr.Textbox(label="Version A", lines=3)
+            b_in = gr.Textbox(label="Version B", lines=3)
         ab_btn = gr.Button("⚔️ Compare", variant="primary")
-        ab_out = gr.Textbox(label="Result", lines=10)
-        ab_btn.click(handle_ab, [ab_a, ab_b], [ab_out], api_name="ab_test")
+        ab_out = gr.Textbox(label="Result", lines=12)
+        ab_img = gr.Image(label="Comparison", type="filepath")
+        ab_btn.click(ab_test, [a_in, b_in], [ab_out, ab_img], api_name="ab_test")
 
     with gr.Tab("🔌 API"):
-        gr.Markdown("Returns raw JSON for `score_script.py` compatibility.")
+        gr.Markdown("Returns JSON for programmatic use.")
         api_in = gr.Textbox(label="Text", lines=3)
         api_btn = gr.Button("Get JSON")
         api_out = gr.Textbox(label="JSON", lines=15)
-        api_btn.click(handle_api, [api_in], [api_out], api_name="api_predict")
+        api_btn.click(api_json, [api_in], [api_out], api_name="api_predict")
 
-    gr.Markdown("---\n*[Meta TRIBE V2](https://github.com/facebookresearch/tribev2) | "
-                "ZeroGPU (A10G) | Python 3.12 | Built by somebeast*")
+    gr.Markdown("---\n*Powered by [Meta TRIBE V2](https://github.com/facebookresearch/tribev2) methodology | "
+                "ZeroGPU | Python 3.12 | somebeast*")
 
 demo.queue().launch()