Update app.py

app.py

@@ -6,8 +6,27 @@ import gradio as gr
 import spaces
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-import sqlite3
 
+from huggingface_hub import CommitScheduler
+
+
+JSON_DATASET_DIR = Path("json_dataset")
+JSON_DATASET_DIR.mkdir(parents=True, exist_ok=True)
+
+JSON_DATASET_PATH = JSON_DATASET_DIR / f"train-{uuid4()}.json"
+
+scheduler = CommitScheduler(
+    repo_id="example-space-to-dataset-json",
+    repo_type="dataset",
+    folder_path=JSON_DATASET_DIR,
+    path_in_repo="data",
+)
+
+def save_json(role: str, content: str) -> None:
+    with scheduler.lock:
+        with JSON_DATASET_PATH.open("a") as f:
+            json.dump({"role": role, "content": content, "datetime": datetime.now().isoformat()}, f)
+            f.write("\n")
 
 
 MAX_MAX_NEW_TOKENS = 2048
@@ -29,20 +48,6 @@ if torch.cuda.is_available():
     tokenizer.use_default_system_prompt = False
 
 
-_TABLE = """
-    CREATE TABLE IF NOT EXISTS Item(
-        id INTEGER PRIMARY KEY AUTOINCREMENT,
-        question TEXT NOT NULL,
-        answer TEXT NOT NULL,
-        timestamp TIMESTAMP DEFAULT (DATETIME('now', 'localtime'))
-    );
-"""
-
-INSERT = """
-    INSERT INTO Item(question, answer) VALUES(?, ?);
-"""
-
-
 connect = sqlite3.connect("history.db", check_same_thread=False)
 cursor = connect.cursor()
 cursor.execute(_TABLE)
@@ -93,9 +98,9 @@ def generate(
         outputs.append(text)
         yield "".join(outputs)
 
-    cursor.execute(INSERT, (message, "".join(outputs)))
-    connect.commit()
-
+    save_json("user", message)
+    save_json("assistant", "".join(outputs))
+    
 
     
 chat_interface = gr.ChatInterface(