Update app.py

app.py

@@ -3,9 +3,51 @@ import transformers
 import torch
 from peft import PeftModel
 import os
+import csv
+import huggingface_hub
+from huggingface_hub import Repository, hf_hub_download, upload_file
+from datetime import datetime
 
+DATASET_REPO_URL = "https://huggingface.co/datasets/JerniganLab/chat-data"
+DATASET_REPO_ID = "JerniganLab/chat-data"
+DATA_FILENAME = "data.csv"
+DATA_FILE = os.path.join("data", DATA_FILENAME)
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
+
+HF_TOKEN = os.environ.get("HF_TOKEN")
+
+# overriding/appending to the gradio template
+SCRIPT = """
+<script>
+if (!window.hasBeenRun) {
+    window.hasBeenRun = true;
+    console.log("should only happen once");
+    document.querySelector("button.submit").click();
+}
+</script>
+"""
+with open(os.path.join(gr.routes.STATIC_TEMPLATE_LIB, "frontend", "index.html"), "a") as f:
+    f.write(SCRIPT)
+
+try:
+    hf_hub_download(
+        repo_id=DATASET_REPO_ID,
+        filename=DATA_FILENAME,
+        cache_dir=DATA_DIRNAME,
+        repo_type='dataset',
+        force_filename=DATA_FILENAME
+    )
+except:
+    print("file not found")
+
+repo = Repository(
+    local_dir="data", clone_from=DATASET_REPO_URL, use_auth_token=HF_TOKEN
+)
+
+
+
+
 model_id = "JerniganLab/interviews-and-qa"
 base_model = "meta-llama/Meta-Llama-3-8B-Instruct"
 
@@ -22,6 +64,17 @@ pipeline = transformers.pipeline(
 
 pipeline.model = PeftModel.from_pretrained(llama_model, model_id)
 
+def store_message(message: str, system_prompt: str, response: str):
+    if response and message:
+        with open(DATA_FILE, "a") as csvfile:
+            writer = csv.DictWriter(csvfile, fieldnames=["message","system_prompt","response","time"])
+            writer.writerow(
+                {"message": message, "system_prompt": system_prompt, "response": response, "time": str(datetime.now())}
+            )
+        commit_url = repo.push_to_hub()
+    # return generate_html()
+
+
 # def chat_function(message, history, system_prompt, max_new_tokens, temperature):
 #     messages = [{"role":"system","content":system_prompt},
 #                 {"role":"user", "content":message}]
@@ -59,6 +112,7 @@ def chat_function(message, history, max_new_tokens, temperature):
         do_sample = True,
         temperature = temperature + 0.1,
         top_p = 0.9,)
+    store_message(message, system_prompt, outputs[0]["generated_text"][len(prompt):])
     return outputs[0]["generated_text"][len(prompt):]
 
 """
@@ -82,10 +136,12 @@ demo = gr.ChatInterface(
     additional_inputs=[
         gr.Slider(100,4000, label="Max New Tokens"),
         gr.Slider(0,1, label="Temperature")
-    ]
+    ], 
+    type="messages",
+    save_history=True,
     )
 
 
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()