Update app.py

app.py

@@ -1,64 +1,84 @@
+import json
 from huggingface_hub import InferenceClient
 import gradio as gr
-import json
+import random
+
+API_URL = "https://api-inference.huggingface.co/models/"
 
 client = InferenceClient(
     "mistralai/Mistral-7B-Instruct-v0.1"
 )
 
-DATABASE_PATH = "database.json"
-
-def load_database():
-    try:
-        with open(DATABASE_PATH, "r") as file:
-            return json.load(file)
-    except FileNotFoundError:
-        return {}
-
-def save_database(database):
-    with open(DATABASE_PATH, "w") as file:
-        json.dump(database, file)
-
 def format_prompt(message, history):
-    prompt = "<s>"
+    prompt = "You're a helpful assistant."
     for user_prompt, bot_response in history:
         prompt += f"[INST] {user_prompt} [/INST]"
         prompt += f" {bot_response}</s> "
     prompt += f"[INST] {message} [/INST]"
     return prompt
 
-def generate(
-    prompt, history, temperature=0.9, max_new_tokens=2000, top_p=0.9, repetition_penalty=1.2,
-):
-    database = load_database()  # Load the database
+def generate(prompt, history, temperature=0.9, max_new_tokens=2048, top_p=0.95, repetition_penalty=1.0):
     temperature = float(temperature)
     if temperature < 1e-2:
         temperature = 1e-2
     top_p = float(top_p)
 
+    generate_kwargs = dict(
+        temperature=temperature,
+        max_new_tokens=max_new_tokens,
+        top_p=top_p,
+        repetition_penalty=repetition_penalty,
+        do_sample=True,
+        seed=random.randint(0, 10**7),
+    )
+
     formatted_prompt = format_prompt(prompt, history)
-    if formatted_prompt in database:
-        response = database[formatted_prompt]
-    else:
-        response = client.text_generation(formatted_prompt, details=True, return_full_text=False)
-        response_text = response.generated_tokens[0].text
-        database[formatted_prompt] = response_text
-        save_database(database)  # Save the updated database
 
-    yield response_text
+    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
+    output = ""
 
-css = """
-  #mkd {
-    height: 500px; 
-    overflow: auto; 
-    border: 1px solid #ccc; 
-  }
-"""
+    for response in stream:
+        output += response.token.text
+        yield output
+    return output
 
-with gr.Blocks(css=css) as demo:
-    gr.ChatInterface(
-        generate,
-        examples=[["What is the secret to life?"], ["Write me a recipe for pancakes."], ["Write a short story about Paris."]]
-    )
+def load_database():
+    try:
+        # Attempt to load the database from JSON
+        with open("database.json", "r", encoding="utf-8") as f:
+            return json.load(f)
+    except (FileNotFoundError, json.JSONDecodeError):
+        # Handle potential errors gracefully
+        print("Error loading database: File not found or invalid format. Creating an empty database.")
+        return []  # Return an empty list if database loading fails
+
+def save_database(data):
+    try:
+        # Save the updated database to JSON
+        with open("database.json", "w", encoding="utf-8") as f:
+            json.dump(data, f, indent=4)
+    except (IOError, json.JSONEncodeError):
+        # Handle potential errors gracefully
+        print("Error saving database: Encountered an issue while saving.")
+
+def chat_interface(message):
+    database = load_database()
+
+    # Check if the question already exists in the database
+    if (message, None) not in database:
+        # If not, generate a response and add it to the database
+        response = generate(message, history=[])
+        database.append((message, response))
+        save_database(database)
+    else:
+        # If it does, retrieve the stored response
+        _, stored_response = next(item for item in database if item[0] == message)
+        response = stored_response
+
+    return response
+
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    demo.register("message", gr.Textbox(label="Your question"))
+    demo.register("response", gr.Textbox(label="Assistant's response"))
 
-demo.launch(debug=True)
+    demo.launch(fn=chat_interface, inputs=["message"], outputs=["response"])