Update app.py

app.py

@@ -1,280 +1,193 @@
-# app.py
-
-import uuid
 import json
-import random
 import traceback
 
-import gradio as gr
-import numpy as np
 import torch
-from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
-from sentence_transformers import SentenceTransformer, util
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, LogitsProcessorList
+from trl import PPOTrainer, PPOConfig
+import gradio as gr
 
 # -----------------------------------------------------------------------------
-# 1. Model & Embedding Initialization
+# 1. Helpers
 # -----------------------------------------------------------------------------
 
-# LLM pipeline cache
-LLM_CACHE = None
-def get_llm(model_name="google/flan-t5-small"):
-    global LLM_CACHE
-    if LLM_CACHE is None:
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
-        model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-        LLM_CACHE = pipeline(
-            "text2text-generation",
-            model=model,
-            tokenizer=tokenizer,
-            device_map="auto" if torch.cuda.is_available() else None,
-            max_length=128,
-        )
-    return LLM_CACHE
-
-# Embedding model for memory & retrieval
-EMB_MODEL = SentenceTransformer("all-MiniLM-L6-v2")
+def make_json_serializable(obj):
+    """
+    Recursively convert any torch.Tensor in obj to Python lists.
+    """
+    if isinstance(obj, torch.Tensor):
+        return obj.cpu().tolist()
+    elif isinstance(obj, dict):
+        return {k: make_json_serializable(v) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [make_json_serializable(v) for v in obj]
+    return obj
+
+def safe_json_dumps(data):
+    """
+    Dump JSON with our converter to avoid Tensor serialization errors.
+    """
+    return json.dumps(
+        make_json_serializable(data),
+        indent=2,
+        ensure_ascii=False
+    )
 
 # -----------------------------------------------------------------------------
-# 2. Trait & Session Definitions
+# 2. Load Models and Initialize PPO Agent
 # -----------------------------------------------------------------------------
 
-BASE_TRAITS = [
-    "Guilt-Proneness", "Anxiety", "Aggression", "Callousness",
-    "Depression", "Grandiosity", "Manipulativeness", "Narcissism",
-    "Impulsivity", "Risk-Taking", "Responsibility", "Empathy",
-    "Conscientiousness"
-]
-EXTENDED_TRAITS = BASE_TRAITS + ["Resilience", "Adaptability"]
+MODEL_NAME = "google/flan-t5-base"
 
-# In‐memory session store
-sessions = {}  # session_id -> session data
-current_session = None
+# Core seq2seq model & tokenizer
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
 
-# -----------------------------------------------------------------------------
-# 3. Question Generation & Scoring
-# -----------------------------------------------------------------------------
+# PPO configuration
+ppo_config = PPOConfig(
+    model_name=MODEL_NAME,
+    learning_rate=1e-5,
+    batch_size=1,
+    log_with=None  # switch to "wandb" or "tensorboard" if you like
+)
 
-def next_question(traits):
-    """Use LLM to generate a probing question for a random trait."""
-    llm = get_llm()
-    trait = random.choice(traits)
-    prompt = f"Generate an open‐ended question to probe the trait '{trait}'."
-    try:
-        q = llm(prompt, do_sample=True, temperature=0.7)[0]["generated_text"].strip()
-    except Exception:
-        q = f"Tell me about a time you felt high in {trait}."
-    return q
-
-def score_response(text, traits):
-    """Rate each trait 0–1 based on the text. Returns mean vector."""
-    llm = get_llm()
-    trait_list = "\n".join(f"{i+1}. {t}" for i, t in enumerate(traits))
-    prompt = (
-        f"Rate the following traits 0–1 from this response:\n\n"
-        f"\"{text}\"\n\n{trait_list}\n\nReturn CSV only."
-    )
-    try:
-        raw = llm(prompt, do_sample=False)[0]["generated_text"]
-        values = [float(x) for x in raw.strip().split(",")]
-        return np.array(values)
-    except Exception:
-        # fallback to neutral
-        return np.full(len(traits), 0.5)
+# Wrap FLAN-T5 in a PPO agent
+ppo_trainer = PPOTrainer(
+    config=ppo_config,
+    model=model,
+    tokenizer=tokenizer
+)
 
 # -----------------------------------------------------------------------------
-# 4. Persona & Memory Functions
+# 3. Session State
 # -----------------------------------------------------------------------------
 
-def build_persona(profile, traits):
-    """Convert numeric profile into textual persona descriptor."""
-    lines = [f"{name}: {int(score*100)}/100"
-             for name, score in zip(traits, profile)]
-    return "### Personality Profile ###\n" + "\n".join(lines)
-
-def embed_text(text):
-    return EMB_MODEL.encode(text, convert_to_tensor=True)
-
-def retrieve_memories(query, memory_store, k=3):
-    """Return top‐k most similar memories to the query."""
-    if not memory_store:
-        return []
-    q_emb = embed_text(query)
-    embs = torch.stack([m["emb"] for m in memory_store])
-    sims = util.pytorch_cos_sim(q_emb, embs)[0]
-    topk = sims.topk(min(k, len(sims)))
-    return [memory_store[i]["text"] for i in topk.indices]
+current_session = {
+    "dialog": []  # each entry: {"user": str, "bot": str, "reward": float or None}
+}
 
 # -----------------------------------------------------------------------------
-# 5. Placeholder: Fine‐Tuning & RL Hooks
+# 4. Core Callback Functions
 # -----------------------------------------------------------------------------
 
-def train_lora_agent(session):
+def reset_session():
     """
-    Placeholder: fine‐tune or LoRA‐adapt a small model
-    on (question, answer) pairs in session["qa_pairs"].
+    Clear the conversation and return an empty chat history.
     """
-    # TODO: integrate peft, LoRA, or Hugging Face Trainer here.
-    pass
+    global current_session
+    current_session = {"dialog": []}
+    return []
 
-def rl_finetune_agent(agent, session):
+def chat_with_agent(user_input: str):
     """
-    Placeholder: reinforce agent reward based on trait‐derived metrics.
-    e.g., higher reward for 'manipulativeness' if persuasion detected.
+    Generate the model's reply, append to session, and return full chat history.
     """
-    # TODO: hook into PPO or other RL frameworks.
-    pass
-
-# -----------------------------------------------------------------------------
-# 6. Evaluation Metrics
-# -----------------------------------------------------------------------------
-
-sentiment_analyzer = pipeline("sentiment-analysis")
-def evaluate_response_style(text):
-    """Return sentiment & basic style metrics."""
-    sent = sentiment_analyzer(text)[0]
-    pronouns = sum(text.lower().count(p) for p in [" i ", " me ", " my "])
-    formality = ("formal" if "you" in text.lower() else "casual")
-    return {
-        "sentiment": sent,
-        "pronouns_used": pronouns,
-        "formality": formality
-    }
-
-# -----------------------------------------------------------------------------
-# 7. Gradio App Logic
-# -----------------------------------------------------------------------------
-
-def start_new_session(num_qs):
     global current_session
-    sid = str(uuid.uuid4())
-    sessions[sid] = {
-        "qa_pairs": [],
-        "profile": None,
-        "persona": None,
-        "memory": []
-    }
-    current_session = sid
-    first_q = next_question(EXTENDED_TRAITS)
-    return sid, first_q
-
-def load_session(json_str, num_qs):
+    try:
+        # Tokenize user prompt and generate
+        inputs = tokenizer(user_input, return_tensors="pt").input_ids
+        outputs = model.generate(
+            inputs,
+            max_new_tokens=128,
+            do_sample=True,
+            top_p=0.9,
+            temperature=0.8
+        )
+        bot_reply = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+        # Store in session
+        current_session["dialog"].append({
+            "user": user_input,
+            "bot": bot_reply,
+            "reward": None
+        })
+
+        # Prepare for Gradio Chatbot: list of (user, bot)
+        history = [
+            (turn["user"], turn["bot"])
+            for turn in current_session["dialog"]
+        ]
+        return history
+    except Exception as e:
+        print("🔥 Error in chat_with_agent:", e)
+        traceback.print_exc()
+        # On failure, leave session untouched
+        return [("Error:", "Failed to generate reply. Check logs.")]
+
+def rate_and_train(rating: float):
+    """
+    Take the last bot reply’s rating, run a PPO step, and return serialized session.
+    """
     global current_session
     try:
-        data = json.loads(json_str)
-        sid = str(uuid.uuid4())
-        sessions[sid] = data
-        current_session = sid
-        # resume with next question
-        return sid, next_question(EXTENDED_TRAITS)
-    except Exception:
-        return None, "Failed to load session JSON."
-
-def submit_answer(answer, num_qs):
-    sess = sessions[current_session]
-    sess["qa_pairs"].append(answer)
-    sess["memory"].append({
-        "text": answer,
-        "emb": embed_text(answer)
-    })
-    if len(sess["qa_pairs"]) < num_qs:
-        return next_question(EXTENDED_TRAITS), None, None
-    # finalize profile
-    all_scores = np.vstack([
-        score_response(txt, EXTENDED_TRAITS) for txt in sess["qa_pairs"]
-    ])
-    final_profile = all_scores.mean(axis=0)
-    sess["profile"] = final_profile.tolist()
-    sess["persona"] = build_persona(final_profile, EXTENDED_TRAITS)
-    # optionally train or RL‐tune
-    train_lora_agent(sess)
-    rl_finetune_agent(None, sess)
-    return None, sess["persona"], json.dumps(sess, indent=2)
-
-def chat_with_agent(user_msg, drift=False):
-    sess = sessions[current_session]
-    history = sess.get("history", [])
-    history.append(("User", user_msg))
-    # retrieve memories
-    mems = retrieve_memories(user_msg, sess["memory"])
-    persona = sess["persona"]
-    llm = get_llm()
-    prompt = (
-        f"{persona}\n\n"
-        f"Relevant Memories:\n" + "\n".join(f"- {m}" for m in mems) + "\n\n"
-        "Conversation History:\n" +
-        "\n".join(f"{s}: {t}" for s, t in history) +
-        "\nAgent:"
-    )
-    out = llm(prompt, do_sample=True, temperature=0.8)[0]["generated_text"]
-    reply = out.split("Agent:")[-1].strip()
-    history.append(("Agent", reply))
-    sess["history"] = history
-    # optional persona drift
-    if drift:
-        # tiny random walk on profile
-        prof = np.array(sess["profile"])
-        sess["profile"] = (prof + np.random.normal(0, 0.01, prof.shape)).clip(0,1).tolist()
-        sess["persona"] = build_persona(np.array(sess["profile"]), EXTENDED_TRAITS)
-    # evaluate style
-    style = evaluate_response_style(reply)
-    return reply, json.dumps(style, indent=2)
+        if not current_session["dialog"]:
+            return "No dialog to rate. Chat first."
+
+        # Attach reward
+        last = current_session["dialog"][-1]
+        last["reward"] = float(rating)
+
+        # Prepare for PPO step
+        user_text = last["user"]
+        bot_text = last["bot"]
+
+        # Token IDs for PPO
+        query_ids = tokenizer(user_text, return_tensors="pt").input_ids.squeeze(0)
+        response_ids = tokenizer(bot_text, return_tensors="pt").input_ids.squeeze(0)
+
+        # Run PPO optimization with this single example
+        stats = ppo_trainer.step(
+            [query_ids],
+            [response_ids],
+            [last["reward"]]
+        )
+        print("🚀 PPO step stats:", stats)
+
+        # Return the entire session as JSON
+        return safe_json_dumps(current_session)
+
+    except Exception as e:
+        print("🔥 Error in rate_and_train:", e)
+        traceback.print_exc()
+        return "Failed to apply training step. See logs."
 
 # -----------------------------------------------------------------------------
-# 8. Build Gradio Interface
+# 5. Gradio UI
 # -----------------------------------------------------------------------------
 
 with gr.Blocks() as demo:
-    gr.Markdown("# Session‐Driven Persona Agent")
+    gr.Markdown("## FLAN-T5 Chatbot with On-the-Fly Reinforcement Learning")
+
+    chat_box = gr.Chatbot(label="Chat History")
+    user_input = gr.Textbox(placeholder="Type your message here…", label="You")
+    send_btn = gr.Button("Send")
+    reset_btn = gr.Button("Reset Conversation")
 
     with gr.Row():
-        mode = gr.Radio(
-            ["New Session", "Load Session"], label="Mode", value="New Session"
-        )
-        num_qs = gr.Slider(3, 10, step=1, label="Number of Profiling Questions", value=5)
-    with gr.Row():
-        load_json = gr.Textbox(
-            label="Paste Session JSON (if loading)", lines=4, visible=False
-        )
-        start_btn = gr.Button("Start Profiling")
+        rating = gr.Slider(0, 5, step=1, value=0, label="Rate Last Reply")
+        rate_btn = gr.Button("Apply Rating & Train")
 
-    sid_box = gr.Textbox(label="Session ID", interactive=False)
-    question_out = gr.Textbox(label="Question", interactive=False)
-    answer_in = gr.Textbox(label="Your Answer")
-    next_btn = gr.Button("Submit Answer")
-    persona_out = gr.Textbox(label="Persona Summary", lines=6)
-    export_json = gr.Textbox(label="Exported Session JSON", lines=6)
+    export_json = gr.Textbox(label="Session JSON", lines=10)
 
-    with gr.Row():
-        user_msg = gr.Textbox(label="Chat with Agent")
-        drift_chk = gr.Checkbox(label="Enable Persona Drift", value=False)
-        chat_btn = gr.Button("Send")
-    chat_out = gr.Textbox(label="Agent Reply", lines=4)
-    style_out = gr.Textbox(label="Reply Style Metrics", lines=4)
-
-    # Show or hide load_json
-    mode.change(lambda m: gr.update(visible=(m=="Load Session")), 
-                inputs=mode, outputs=load_json)
-
-    # Start or load session
-    start_btn.click(
-        fn=lambda m, n, js: load_session(js, n) if m=="Load Session" else start_new_session(n),
-        inputs=[mode, num_qs, load_json],
-        outputs=[sid_box, question_out]
+    # Reset chat
+    reset_btn.click(
+        fn=reset_session,
+        inputs=None,
+        outputs=chat_box
     )
 
-    # Submit profiling answer
-    next_btn.click(
-        fn=lambda ans, n: submit_answer(ans, int(n)),
-        inputs=[answer_in, num_qs],
-        outputs=[question_out, persona_out, export_json]
+    # Send user message
+    send_btn.click(
+        fn=chat_with_agent,
+        inputs=user_input,
+        outputs=chat_box
     )
 
-    # Chat interface
-    chat_btn.click(
-        fn=chat_with_agent,
-        inputs=[user_msg, drift_chk],
-        outputs=[chat_out, style_out]
+    # Rate & train
+    rate_btn.click(
+        fn=rate_and_train,
+        inputs=rating,
+        outputs=export_json
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()