Update app.py

app.py

@@ -12,28 +12,24 @@ models = {
         "mistralai/Mistral-7B-Instruct-v0.3",
         device_map="auto",
         torch_dtype=torch.bfloat16,  # Use bfloat16
-        use_auth_token=True,
-        token=hf_token  # Add authentication here
+        token=hf_token  # Use token for authentication
     ),
     "BICORP/Lake-1-Advanced": AutoModelForCausalLM.from_pretrained(
         "BICORP/Lake-1-Advanced",
         device_map="auto",
         torch_dtype=torch.bfloat16,  # Use bfloat16
-        use_auth_token=True,
-        token=hf_token  # Add authentication here
+        token=hf_token  # Use token for authentication
     )
 }
 
 tokenizers = {
     "mistralai/Mistral-7B-Instruct-v0.3": AutoTokenizer.from_pretrained(
         "mistralai/Mistral-7B-Instruct-v0.3",
-        use_auth_token=True,
-        token=hf_token  # Add authentication here
+        token=hf_token  # Use token for authentication
     ),
     "BICORP/Lake-1-Advanced": AutoTokenizer.from_pretrained(
         "BICORP/Lake-1-Advanced",
-        use_auth_token=True,
-        token=hf_token  # Add authentication here
+        token=hf_token  # Use token for authentication
     )
 }
 
@@ -73,7 +69,7 @@ def respond(message, history: list, model_name, preset_name):
     preset = presets[model_name][preset_name]
 
     # Prepare the input for the model
-    input_text = f"{system_messages[model_name]}\n:User   {message}\nAI:"
+    input_text = f"{system_messages[model_name]}\n:User    {message}\nAI:"
     inputs = tokenizer.encode(input_text, return_tensors="pt").to(model.device)
 
     # Generate response
@@ -89,33 +85,43 @@ def respond(message, history: list, model_name, preset_name):
     response = tokenizer.decode(output[0], skip_special_tokens=True)
     return response.split("AI:")[-1].strip()  # Extract the AI's response
 
-def respond_with_pseudonym(message, history, selected_model, selected_preset):
-    try:
-        model_id = next(model[0] for model in model_choices if model[1] == selected_model)
-    except StopIteration:
-        return "Model not found."
-
-    response = respond(message, history, model_id, selected_preset)
-    history.append((message, response))
-    return history
-
-# Gradio interface setup
-def launch_interface():
-    with gr.Blocks() as demo:
-        gr.Markdown("## Chat with AI Models")
-        with gr.Row():
-            model_dropdown = gr.Dropdown(choices=pseudonyms, label="Select Model")
-            preset_dropdown = gr.Dropdown(choices=["Fast", "Normal", "Quality", "Unreal Performance"], label="Select Preset")
-        chatbot = gr.Chatbot()
-        message_input = gr.Textbox(placeholder="Type your message here...")
-        submit_button = gr.Button("Send")
+def respond_with_pseudonym(message, history: list, model_name, preset_name, pseudonym):
+    # Get the correct model and tokenizer
+    model = models[model_name]
+    tokenizer = tokenizers[model_name]
+    preset = presets[model_name][preset_name]
 
-        def submit_message(message, history, selected_model, selected_preset):
-            return respond_with_pseudonym(message, history, selected_model, selected_preset)
+    # Prepare the input for the model with pseudonym
+    input_text = f"{system_messages[model_name]}\n:{pseudonym}    {message}\nAI:"
+    inputs = tokenizer.encode(input_text, return_tensors="pt").to(model.device)
 
-        submit_button.click(submit_message, inputs=[message_input, chatbot, model_dropdown, preset_dropdown], outputs=chatbot)
+    # Generate response
+    with torch.no_grad():
+        output = model.generate(
+            inputs,
+            max_new_tokens=preset["max_new_tokens"],
+            temperature=preset["temperature"],
+            top_p=preset["top_p"]
+        )
 
-    demo.launch(share=True)
+    # Decode the output
+    response = tokenizer.decode(output[0], skip_special_tokens=True)
+    return response.split("AI:")[-1].strip()  # Extract the AI's response
 
-if __name__ == "__main__":
-    launch_interface()
+# Gradio interface setup
+iface = gr.Interface(
+    fn=respond_with_pseudonym,
+    inputs=[
+        gr.inputs.Textbox(label="Message"),
+        gr.inputs.State(),
+        gr.inputs.Dropdown(choices=pseudonyms, label="Model"),
+        gr.inputs.Dropdown(choices=["Fast", "Normal", "Quality", "Unreal Performance"], label="Preset"),
+        gr.inputs.Textbox(label="Pseudonym", default="User ")
+    ],
+    outputs="text",
+    title="AI Chatbot",
+    description="Chat with AI models using your chosen pseudonym."
+)
+
+# Launch the Gradio app
+iface.launch()    launch_interface()