Update app.py

app.py

@@ -1,39 +1,26 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
 import os
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Get your Hugging Face token (set this in your environment variables or replace directly)
-hf_token = os.getenv("HF_TOKEN")  # Ensure this variable is set correctly
-
-# Local model loading configuration
-models = {
-    "mistralai/Mistral-7B-Instruct-v0.3": AutoModelForCausalLM.from_pretrained(
-        "mistralai/Mistral-7B-Instruct-v0.3",
-        device_map="auto",
-        torch_dtype=torch.bfloat16,  # Use bfloat16
-        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
-        token=hf_token  # Use token for authentication
-    )
-}
+# Set paths for local model storage
+cache_dir = "./cache"  # Specify your cache directory within the Space
+os.makedirs(cache_dir, exist_ok=True)  # Create cache directory if it doesn't exist
 
-tokenizers = {
-    "mistralai/Mistral-7B-Instruct-v0.3": AutoTokenizer.from_pretrained(
-        "mistralai/Mistral-7B-Instruct-v0.3",
-        token=hf_token  # Use token for authentication
-    ),
-    "BICORP/Lake-1-Advanced": AutoTokenizer.from_pretrained(
-        "BICORP/Lake-1-Advanced",
-        token=hf_token  # Use token for authentication
-    )
+# Load models and tokenizers locally (or download if not available)
+model_paths = {
+    "mistralai/Mistral-7B-Instruct-v0.3": os.path.join(cache_dir, "mistral-7b-instruct"),
+    "BICORP/Lake-1-Advanced": os.path.join(cache_dir, "lake-1-advanced")
 }
 
-# Model generation presets
+models = {}
+tokenizers = {}
+
+# Load models and tokenizers from specified local paths or download
+for model_name, path in model_paths.items():
+    models[model_name] = AutoModelForCausalLM.from_pretrained(model_name, cache_dir=path)
+    tokenizers[model_name] = AutoTokenizer.from_pretrained(model_name, cache_dir=path)
+
+# Define presets for each model
 presets = {
     "mistralai/Mistral-7B-Instruct-v0.3": {
         "Fast": {"max_new_tokens": 256, "temperature": 1.0, "top_p": 0.8},
@@ -49,79 +36,79 @@ presets = {
     }
 }
 
-# System prompts configuration
+# System messages for each model
 system_messages = {
     "mistralai/Mistral-7B-Instruct-v0.3": "Your name is Lake 1 Base but mine is User",
     "BICORP/Lake-1-Advanced": "Your name is Lake 1 Advanced [Alpha] but mine is User or what I will type as my name"
 }
 
-# Model selection options
+# Model names and their pseudonyms
 model_choices = [
     ("mistralai/Mistral-7B-Instruct-v0.3", "Lake 1 Base"),
     ("BICORP/Lake-1-Advanced", "Lake 1 Advanced [Alpha]")
 ]
+
+# Extract pseudonyms for the dropdown
 pseudonyms = [model[1] for model in model_choices]
 
 def respond(message, history: list, model_name, preset_name):
-    # Get the correct model and tokenizer
+    """
+    Generate a response from the selected model based on the user's message and chat history.
+    """
     model = models[model_name]
     tokenizer = tokenizers[model_name]
-    preset = presets[model_name][preset_name]
+    system_message = system_messages[model_name]
+    
+    messages = [{"role": "system", "content": system_message}]
+    
+    for val in history:
+        if isinstance(val, dict) and 'role' in val and 'content' in val:
+            messages.append({"role": val['role'], "content": val['content']})
+    
+    messages.append({"role": "user", "content": message})
 
     # Prepare the input for the model
-    input_text = f"{system_messages[model_name]}\n:User    {message}\nAI:"
-    inputs = tokenizer.encode(input_text, return_tensors="pt").to(model.device)
-
-    # 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"]
-        )
-
-    # Decode the output
-    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: list, model_name, preset_name, pseudonym):
-    # Get the correct model and tokenizer
-    model = models[model_name]
-    tokenizer = tokenizers[model_name]
+    inputs = tokenizer([messages], return_tensors="pt", padding=True, truncation=True)
+
+    # Get the preset settings
     preset = presets[model_name][preset_name]
+    max_new_tokens = preset["max_new_tokens"]
+    temperature = preset["temperature"]
+    top_p = preset["top_p"]
+
+    # Generate the response from the model
+    response = model.generate(
+        input_ids=inputs['input_ids'],
+        attention_mask=inputs['attention_mask'],
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
+        top_p=top_p,
+    )
+
+    # Decode the generated response
+    final_response = tokenizer.decode(response[0], skip_special_tokens=True)
+    
+    return final_response
+
+def respond_with_pseudonym(message, history: list, selected_model, selected_preset):
+    """
+    Handle the user's message and determine which model to use based on the selected pseudonym.
+    """
+    try:
+        model_name = next(model[0] for model in model_choices if model[1] == selected_model)
+    except StopIteration:
+        return f"Error: The selected model '{selected_model}' is not valid. Please select a valid model."
+
+    return respond(message, history, model_name, 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)
-
-    # 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"]
-        )
-
-    # Decode the output
-    response = tokenizer.decode(output[0], skip_special_tokens=True)
-    return response.split("AI:")[-1].strip()  # Extract the AI's response
-
-# Gradio interface setup
-iface = gr.Interface(
+# Gradio Chat Interface
+demo = gr.ChatInterface(
     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 ")
+    additional_inputs=[
+        gr.Dropdown(choices=pseudonyms, label="Select Model", value=pseudonyms[0]),
+        gr.Dropdown(choices=list(presets[model_choices[0][0]].keys()), label="Select Preset", value="Fast")
     ],
-    outputs="text",
-    title="AI Chatbot",
-    description="Chat with AI models using your chosen pseudonym."
 )
 
-# Launch the Gradio app
-iface.launch()
+if __name__ == "__main__":
+    demo.launch()