add writing to dataset

app.py

@@ -8,42 +8,38 @@ import uuid
 with open("system_prompt.txt", "r") as f:
     SYSTEM_PROMPT = f.read()
 
+# ---- Constants ----
 MODEL_NAME = "HuggingFaceH4/zephyr-7b-beta"
+DATASET_REPO = "frimelle/companion-chat-logs"
+HF_TOKEN = os.environ.get("HF_TOKEN")  # set in Space secrets
 client = InferenceClient(MODEL_NAME)
 
-# ---- Setup logging ----
-LOG_DIR = "chat_logs"
-os.makedirs(LOG_DIR, exist_ok=True)
-session_id = str(uuid.uuid4())
+# ---- Upload to Dataset ----
+def upload_chat_to_dataset(user_message, assistant_message, system_prompt):
+    row = {
+        "timestamp": datetime.now().isoformat(),
+        "session_id": str(uuid.uuid4()),
+        "user": user_message,
+        "assistant": assistant_message,
+        "system_prompt": system_prompt,
+    }
 
-def log_chat(session_id, user_msg, bot_msg):
-    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-    log_path = os.path.join(LOG_DIR, f"{session_id}.txt")
-    with open(log_path, "a", encoding="utf-8") as f:
-        f.write(f"[{timestamp}] User: {user_msg}\n")
-        f.write(f"[{timestamp}] Bot: {bot_msg}\n\n")
+    dataset = Dataset.from_dict({k: [v] for k, v in row.items()})
+    dataset.push_to_hub(DATASET_REPO, private=True, token=HF_TOKEN)
 
-# ---- Respond Function with Logging ----
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
+# ---- Chat Function ----
+def respond(message, history, system_message, max_tokens, temperature, top_p):
     messages = [{"role": "system", "content": system_message}]
 
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
+    for user_msg, bot_msg in history:
+        if user_msg:
+            messages.append({"role": "user", "content": user_msg})
+        if bot_msg:
+            messages.append({"role": "assistant", "content": bot_msg})
 
     messages.append({"role": "user", "content": message})
 
     response = ""
-
     for chunk in client.chat_completion(
         messages,
         max_tokens=max_tokens,
@@ -56,19 +52,19 @@ def respond(
             response += token
             yield response
 
-    # Save full message after stream ends
-    log_chat(session_id, message, response)
+    # Log the final full message to the dataset
+    upload_chat_to_dataset(message, response, system_message)
 
-# ---- Gradio Interface ----
+# ---- Gradio UI ----
 demo = gr.ChatInterface(
-    respond,
+    fn=respond,
     additional_inputs=[
         gr.Textbox(value=SYSTEM_PROMPT, label="System message"),
         gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
         gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
         gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
     ],
-    title="BoundrAI"
+    title="BoundrAI",
 )
 
 if __name__ == "__main__":

requirements.txt

@@ -1 +1,3 @@
-huggingface_hub==0.25.2
+huggingface_hub==0.25.2
+gradio
+datasets