Update app.py

app.py

@@ -4,55 +4,161 @@ from huggingface_hub import login
 import torch
 import os
 
+# Hugging Face token login
 HF_TOKEN = os.getenv("HF_TOKEN")
 login(token=HF_TOKEN)
 
+# Define models
+MODELS = {
+    "atlas-flash-1215": {
+        "name": "🦁 Atlas-Flash 1215",
+        "sizes": {
+            "1.5B": "Spestly/Atlas-Flash-1.5B-Preview",
+        },
+        "emoji": "🦁",
+        "experimental": True,
+        "is_vision": False,
+        "system_prompt_env": "ATLAS_FLASH_1215",
+    },
+    "atlas-pro-0403": {
+        "name": "🏆 Atlas-Pro 0403",
+        "sizes": {
+            "1.5B": "Spestly/Atlas-Pro-1.5B-Preview",
+        },
+        "emoji": "🏆",
+        "experimental": True,
+        "is_vision": False,
+        "system_prompt_env": "ATLAS_PRO_0403",
+    },
 
-model_name = "Spestly/Atlas-Pro-1.5B-Preview"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float32, low_cpu_mem_usage=True)
+}
 
+# Load default model
+default_model_key = "atlas-pro-0403"
+default_size = "1.5B"
+default_model = MODELS[default_model_key]["sizes"][default_size]
 
-model.eval()
-
-def generate_response(message, history):
-    instruction = (
-        "You are an LLM called Atlas. You are finetuned by Aayan Mishra. You are NOT trained by Anthropic. "
-        "You are a Qwen 2.5 fine-tune. Your purpose is the help the user accomplish their request to the best of your abilities. "
-        "Below is an instruction that describes a task. Answer it clearly and concisely.\n\n"
-        f"### Instruction:\n{message}\n\n### Response:"
+def load_model(model_name):
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name, 
+        torch_dtype=torch.float32, 
+        low_cpu_mem_usage=True
     )
+    model.eval()
+    return tokenizer, model
+
+tokenizer, model = load_model(default_model)
+
+# Generate response function
+def generate_response(message, image, history, model_key, model_size, temperature, top_p, max_new_tokens):
+    global tokenizer, model
+    # Load the selected model
+    selected_model = MODELS[model_key]["sizes"][model_size]
+    if selected_model != default_model:
+        tokenizer, model = load_model(selected_model)
     
+    # Get the system prompt from the environment
+    system_prompt_env = MODELS[model_key]["system_prompt_env"]
+    system_prompt = os.getenv(system_prompt_env, "You are an advanced AI system. Help the user as best as you can.")
+    
+    # Construct instruction
+    if MODELS[model_key]["is_vision"]:
+        # If a vision model, include the image information
+        image_info = "An image has been provided as input."
+        instruction = (
+            f"{system_prompt}\n\n"
+            f"### Instruction:\n{message}\n{image_info}\n\n### Response:"
+        )
+    else:
+        # For non-vision models
+        instruction = (
+            f"{system_prompt}\n\n"
+            f"### Instruction:\n{message}\n\n### Response:"
+        )
+    
+    # Tokenize input
     inputs = tokenizer(instruction, return_tensors="pt")
     
+    # Generate response
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=1000,
+            max_new_tokens=max_new_tokens,
             num_return_sequences=1,
-            temperature=0.7,
-            top_p=0.9,
+            temperature=temperature,
+            top_p=top_p,
             do_sample=True
         )
     
+    # Decode response
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     response = response.split("### Response:")[-1].strip()
-    
     return response
 
-iface = gr.ChatInterface(
-    generate_response,
-    chatbot=gr.Chatbot(height=600, type="messages"),
-    textbox=gr.Textbox(placeholder="Type your message here...", container=False, scale=7),
-    title="🦁 Atlas-Pro",
-    description="Chat with Alas-Pro",
-    theme="soft",
-    examples=[
-        "Can you give me a good salsa recipe?",
-        "Write an engaging two-line horror story.",
-        "What is the capital of Australia?",
-    ],
-    type="messages"
-)
-
-iface.launch()
+# User interface
+def create_interface():
+    # Define input components
+    message_input = gr.Textbox(label="Message", placeholder="Type your message here...")
+    model_key_selector = gr.Dropdown(
+        label="Model",
+        choices=list(MODELS.keys()),
+        value=default_model_key
+    )
+    model_size_selector = gr.Dropdown(
+        label="Model Size",
+        choices=list(MODELS[default_model_key]["sizes"].keys()),
+        value=default_size
+    )
+    temperature_slider = gr.Slider(label="Temperature", minimum=0.1, maximum=1.0, value=0.7, step=0.1)
+    top_p_slider = gr.Slider(label="Top-p", minimum=0.1, maximum=1.0, value=0.9, step=0.1)
+    max_tokens_slider = gr.Slider(label="Max New Tokens", minimum=50, maximum=2000, value=1000, step=50)
+    image_input = gr.Image(label="Upload Image (if applicable)", type="filepath", visible=False)
+
+    # Function to toggle visibility of image input
+    def toggle_image_input(model_key):
+        return MODELS[model_key]["is_vision"]
+
+    # Output components
+    chat_output = gr.Chatbot(label="Chatbot")
+
+    # Function to process inputs and generate output
+    def process_inputs(message, image, model_key, model_size, temperature, top_p, max_new_tokens, history=[]):
+        response = generate_response(
+            message=message,
+            image=image,
+            history=history,
+            model_key=model_key,
+            model_size=model_size,
+            temperature=temperature,
+            top_p=top_p,
+            max_new_tokens=max_new_tokens
+        )
+        history.append((message, response))
+        return history
+
+    # Interface layout
+    iface = gr.Interface(
+        fn=process_inputs,
+        inputs=[
+            message_input,
+            image_input,
+            model_key_selector,
+            model_size_selector,
+            temperature_slider,
+            top_p_slider,
+            max_tokens_slider
+        ],
+        outputs=chat_output,
+        title="🌟 Atlas-Pro/Flash/Vision Interface",
+        description="Interact with multiple models like Atlas-Pro, Atlas-Flash, and AtlasV-Pro (Comming Soon!). Upload images for vision models!",
+        theme="soft",
+        live=True
+    )
+
+    # Add event to toggle image input visibility
+    iface.input_components[1].set_visibility(toggle_image_input(model_key_selector.value))
+    return iface
+
+# Launch the app
+create_interface().launch()