add logging to dataset

app.py

@@ -1,18 +1,53 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+from huggingface_hub import InferenceClient, HfApi
+from datetime import datetime
+import uuid
+import os
+import json
 
-# ---- System prompt and model ----
+# ---- Configuration ----
+MODEL_NAME = "HuggingFaceH4/zephyr-7b-beta"
+DATASET_REPO = "frimelle/companion-chat-logs"
+HF_TOKEN = os.environ.get("HF_TOKEN")
+
+# ---- Load system prompt ----
 with open("system_prompt.txt", "r") as f:
     SYSTEM_PROMPT = f.read()
 
-MODEL_NAME = "HuggingFaceH4/zephyr-7b-beta"
 client = InferenceClient(MODEL_NAME)
+api = HfApi()
 
-# ---- Fixed parameters ----
+# ---- Fixed generation parameters ----
 MAX_TOKENS = 512
 TEMPERATURE = 0.7
 TOP_P = 0.95
 
+# ---- Logging function ----
+def log_to_dataset(user_message, assistant_message):
+    row = {
+        "timestamp": datetime.now().isoformat(),
+        "session_id": str(uuid.uuid4()),
+        "user": user_message,
+        "assistant": assistant_message,
+        "system_prompt": SYSTEM_PROMPT,
+    }
+
+    # Save as a JSONL file and upload to the dataset repo
+    filename = f"log_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:8]}.jsonl"
+    with open(filename, "w", encoding="utf-8") as f:
+        f.write(json.dumps(row) + "\n")
+
+    api.upload_file(
+        path_or_fileobj=filename,
+        path_in_repo=filename,
+        repo_id=DATASET_REPO,
+        repo_type="dataset",
+        token=HF_TOKEN
+    )
+
+    os.remove(filename)
+
+# ---- Chatbot function ----
 def respond(message, history):
     messages = [{"role": "system", "content": SYSTEM_PROMPT}]
     
@@ -37,10 +72,13 @@ def respond(message, history):
             response += token
             yield response
 
-# ---- Simple chat interface ----
+    # Log after full response is generated
+    log_to_dataset(message, response)
+
+# ---- Gradio Interface ----
 demo = gr.ChatInterface(
     respond,
-    title="BoundrAI",
+    title="Zephyr Chatbot",
 )
 
 if __name__ == "__main__":