Update app.py

app.py

@@ -1,13 +1,12 @@
 from __future__ import annotations
 
-# -- standard lib
 import json
 import os
 import time
 import uuid
 from threading import Thread
+from typing import List, Dict
 
-# -- third-party deps (declared in requirements.txt of the Space)
 import gradio as gr
 from gradio_modal import Modal
 from transformers import (
@@ -15,30 +14,32 @@ from transformers import (
     AutoTokenizer,
     TextIteratorStreamer,
 )
-from datasets import Dataset, load_dataset, concatenate_datasets, DownloadMode
+from datasets import (
+    Dataset,
+    load_dataset,
+    concatenate_datasets,
+    DownloadMode,
+)
 from huggingface_hub import HfApi, login
 import spaces
 
-# ──────────────────────────── model & constants ─────────────────────────────
+# ─────────────────────────── model & constants ────────────────────────────
 checkpoint = "marin-community/marin-8b-instruct"
-device = "cuda"  # the Space runner gives us a GPU
+device = "cuda"
 
-# download 🔥
 tokenizer = AutoTokenizer.from_pretrained(checkpoint)
 model = AutoModelForCausalLM.from_pretrained(checkpoint).to(device)
 
-# feedback dataset details
-DATASET_REPO = "WillHeld/model-feedback"  # <-- change to your namespace if needed
+DATASET_REPO = "WillHeld/model-feedback"  # change if forking
 DATA_DIR = "./feedback_data"
 DATA_FILE = "feedback.jsonl"
 os.makedirs(DATA_DIR, exist_ok=True)
 
-# ──────────────────────────── helpers ───────────────────────────────────────
+# ─────────────────────────── helper functions ─────────────────────────────
 
-def save_feedback_locally(conversation: list[dict[str, str]],
+def save_feedback_locally(conversation: List[Dict[str, str]],
                           satisfaction: str,
-                          feedback_text: str) -> str:
-    """Append a single feedback record to a JSONL file and return its UUID."""
+                          feedback_text: str) -> None:
     record = {
         "id": str(uuid.uuid4()),
         "timestamp": time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),
@@ -46,39 +47,25 @@ def save_feedback_locally(conversation: list[dict[str, str]],
         "satisfaction": satisfaction,
         "feedback": feedback_text,
     }
-    fp = os.path.join(DATA_DIR, DATA_FILE)
-    with open(fp, "a", encoding="utf-8") as f:
-        f.write(json.dumps(record, ensure_ascii=False) + "\n")
-    return record["id"]
-
+    with open(os.path.join(DATA_DIR, DATA_FILE), "a", encoding="utf-8") as fp:
+        fp.write(json.dumps(record, ensure_ascii=False) + "\n")
 
-def push_feedback_to_hub(hf_token: str | None = None) -> bool:  # noqa: C901
-    """Merge freshly collected feedback with what’s already on the Hub.
-
-    Steps
-    -----
-    1.  Authenticate with `hf_token` (fall back to $HF_TOKEN env).
-    2.  Load *local* feedback just written in `feedback.jsonl`.
-    3.  Pull existing remote split (if any); concat & `unique("id")`.
-    4.  Push the merged dataset back. Never deletes remote shards ⇒ safe.
-    """
 
+def push_feedback_to_hub(hf_token: str | None = None) -> bool:
     hf_token = hf_token or os.getenv("HF_TOKEN")
     if not hf_token:
         print("❌  No HF token — skipping Hub push.")
         return False
     login(token=hf_token)
 
-    fp = os.path.join(DATA_DIR, DATA_FILE)
-    if not os.path.exists(fp):
-        print("❌  Local feedback file missing; nothing to push.")
+    local_path = os.path.join(DATA_DIR, DATA_FILE)
+    if not os.path.exists(local_path):
+        print("❌  No local feedback to push.")
         return False
 
-    # local rows → Dataset
-    with open(fp, encoding="utf-8") as f:
-        local_ds = Dataset.from_list([json.loads(l) for l in f])
+    with open(local_path, encoding="utf-8") as fp:
+        local_ds = Dataset.from_list([json.loads(l) for l in fp])
 
-    # try to pull remote
     try:
         remote_ds = load_dataset(
             DATASET_REPO,
@@ -88,10 +75,8 @@ def push_feedback_to_hub(hf_token: str | None = None) -> bool:  # noqa: C901
         )
         merged = concatenate_datasets([remote_ds, local_ds]).unique("id")
     except FileNotFoundError:
-        # repo exists but empty
         merged = local_ds
     except Exception:
-        # repo may not exist yet – create & start fresh
         HfApi(token=hf_token).create_repo(
             repo_id=DATASET_REPO, repo_type="dataset", private=True
         )
@@ -102,25 +87,23 @@ def push_feedback_to_hub(hf_token: str | None = None) -> bool:  # noqa: C901
         private=True,
         commit_message=f"Add {len(local_ds)} new feedback entries",
     )
-    print(
-        f"✅  Pushed {len(local_ds)} rows; dataset now has {len(merged)} total.")
-    # (optional) clear local file once synced
-    # os.remove(fp)
+    print(f"✅  Pushed {len(local_ds)} rows; dataset now has {len(merged)} total.")
     return True
 
-# ──────────────────────────── chat backend ─────────────────────────────────
+# ─────────────────────────── chat backend ────────────────────────────────
 
 @spaces.GPU(duration=120)
 def generate_response(message: str,
-                      history: list[dict[str, str]],
+                      history: List[Dict[str, str]],
+                      conversation_state: List[Dict[str, str]],
                       temperature: float,
                       top_p: float):
-    """Streaming generator used by the Gradio ChatInterface."""
+    """Yields assistant text only; conversation_state is updated in‑place."""
 
-    # 1) add user message to history
+    # sync state
     history.append({"role": "user", "content": message})
+    conversation_state[:] = history  # keep external state in sync
 
-    # 2) build model input via chat template
     prompt = tokenizer.apply_chat_template(history, tokenize=False,
                                            add_generation_prompt=True)
     input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
@@ -137,79 +120,61 @@ def generate_response(message: str,
         streamer=streamer,
     )
 
-    # run on a worker thread so we can yield tokens live
     Thread(target=model.generate, kwargs=gen_kwargs).start()
 
     partial = ""
     for token in streamer:
         partial += token
-        yield partial, history  # 1st out = msg, 2nd out = state
+        yield partial  # only the assistant text gets streamed
 
-    # once finished, commit assistant reply to history
     history.append({"role": "assistant", "content": partial})
-    yield partial, history
+    conversation_state[:] = history
+    # (no final yield; generator simply ends)
 
-# ──────────────────────────── feedback handler ─────────────────────────────
+# ─────────────────────────── feedback handler ────────────────────────────
 
-def submit_feedback(conversation_state: list[dict[str, str]],
+def submit_feedback(conversation_state: List[Dict[str, str]],
                     satisfaction: str,
                     feedback_text: str):
-    """Callback for the *Submit Research Feedback* button."""
     save_feedback_locally(conversation_state, satisfaction, feedback_text)
-    pushed = push_feedback_to_hub()
-    if pushed:
+    if push_feedback_to_hub():
         return "✅  Thanks! Your feedback is safely stored."
     return "⚠️  Saved locally; Hub push failed. Check server logs."
 
-# ──────────────────────────── UI layout ────────────────────────────────────
+# ─────────────────────────── UI layout ───────────────────────────────────
 
-with gr.Blocks(title="Marin-8B Research Preview") as demo:
-    # state object to surface chat history to the feedback form
+with gr.Blocks(title="Marin‑8B Research Preview") as demo:
     conversation_state = gr.State([])
 
     with gr.Row():
-        # ——— Chat column ———
         with gr.Column(scale=3):
             chatbot = gr.ChatInterface(
                 fn=generate_response,
-                additional_inputs=[conversation_state,  # keeps state in sync
+                additional_inputs=[conversation_state,
                                    gr.Slider(0.1, 2.0, value=0.7, step=0.1,
                                              label="Temperature"),
                                    gr.Slider(0.1, 1.0, value=0.9, step=0.05,
-                                             label="Top-P")],
-                additional_outputs=[conversation_state],
+                                             label="Top‑P")],
                 type="messages",
             )
-
-        # ——— Sidebar column ———
         with gr.Column(scale=1):
             report_btn = gr.Button("Share Feedback", variant="primary")
 
-    # feedback modal (hidden by default)
     with Modal(visible=False) as fb_modal:
         gr.Markdown("## Research Preview Feedback")
-        gr.Markdown("We appreciate your help improving Marin-8B! ✨")
-
+        gr.Markdown("We appreciate your help improving Marin‑8B! ✨")
         sat_radio = gr.Radio([
             "Very satisfied", "Satisfied", "Neutral",
             "Unsatisfied", "Very unsatisfied"],
-            label="Overall experience",
-            value="Neutral",
-        )
+            label="Overall experience", value="Neutral")
         fb_text = gr.Textbox(lines=6, label="Comments / suggestions")
         send_btn = gr.Button("Submit", variant="primary")
         status_box = gr.Textbox(label="Status", interactive=False)
 
-    # interactions
-    # open the modal without custom JS – use Modal update
-    report_btn.click(lambda: Modal(visible=True), None, fb_modal)
-
-    send_btn.click(
-        submit_feedback,
-        inputs=[conversation_state, sat_radio, fb_text],
-        outputs=status_box,
-    )
+    report_btn.click(lambda: Modal.update(visible=True), None, fb_modal)
+    send_btn.click(submit_feedback,
+                   inputs=[conversation_state, sat_radio, fb_text],
+                   outputs=status_box)
 
-# ──────────────────────────── run! ─────────────────────────────────────────
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()