Upload 2 files

app.py

@@ -1,8 +1,20 @@
 import os
+
 import gradio as gr
 from text_generation import Client
 from conversation import get_default_conv_template
 from transformers import AutoTokenizer
+from pymongo import MongoClient
+
+DB_NAME = os.getenv("MONGO_DBNAME", "taiwan-llm")
+USER = os.getenv("MONGO_USER")
+PASSWORD = os.getenv("MONGO_PASSWORD")
+
+uri = f"mongodb+srv://{USER}:{PASSWORD}@{DB_NAME}.kvwjiok.mongodb.net/?retryWrites=true&w=majority"
+mongo_client = MongoClient(uri)
+db = mongo_client[DB_NAME]
+conversations_collection = db['conversations']
+
 DESCRIPTION = """
 # Language Models for Taiwanese Culture
 
@@ -64,7 +76,8 @@ DEFAULT_MAX_NEW_TOKENS = 1024
 
 max_prompt_length = 4096 - MAX_MAX_NEW_TOKENS - 10
 
-tokenizer = AutoTokenizer.from_pretrained("yentinglin/Taiwan-LLaMa-v1.0")
+model_name = "yentinglin/Taiwan-LLaMa-v1.0"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 with gr.Blocks() as demo:
     gr.Markdown(DESCRIPTION)
@@ -152,6 +165,18 @@ with gr.Blocks() as demo:
                 history[-1][1] += character
                 yield history
 
+        # After generating the response, store the conversation history in MongoDB
+        conversation_document = {
+            "model_name": model_name,
+            "history": history,
+            "system_prompt": system_prompt,
+            "max_new_tokens": max_new_tokens,
+            "temperature": temperature,
+            "top_p": top_p,
+            "top_k": top_k,
+        }
+        conversations_collection.insert_one(conversation_document)
+
     msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
         fn=bot,
         inputs=[
@@ -238,4 +263,4 @@ with gr.Blocks() as demo:
     gr.Markdown(LICENSE)
 
 demo.queue(concurrency_count=4, max_size=128)
-demo.launch()
+demo.launch()

requirements.txt

@@ -1,2 +1,3 @@
 text-generation==0.6.0
-transformers==4.31.0
+transformers==4.31.0
+pymongo==4.4.1