Update app.py

app.py

@@ -30,7 +30,6 @@ MODELS = {
         "is_vision": False,
         "system_prompt_env": "ATLAS_PRO_0403",
     },
-
 }
 
 # Load default model
@@ -53,34 +52,21 @@ 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:"
-        )
+        instruction = f"{system_prompt}\n\n### 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:"
-        )
+        instruction = f"{system_prompt}\n\n### Instruction:\n{message}\n\n### Response:"
     
-    # Tokenize input
     inputs = tokenizer(instruction, return_tensors="pt")
     
-    # Generate response
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
@@ -91,74 +77,61 @@ def generate_response(message, image, history, model_key, model_size, temperatur
             do_sample=True
         )
     
-    # Decode response
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     response = response.split("### Response:")[-1].strip()
     return response
 
-# 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)
+    with gr.Blocks(title="🌟 Atlas-Pro/Flash/Vision Interface", theme="soft") as iface:
+        gr.Markdown("Interact with multiple models like Atlas-Pro, Atlas-Flash, and AtlasV-Flash (Coming Soon!). Upload images for vision models!")
 
-    # Function to toggle visibility of image input
-    def toggle_image_input(model_key):
-        return MODELS[model_key]["is_vision"]
+        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
+        )
+        image_input = gr.Image(label="Upload Image (if applicable)", type="filepath", visible=False)
+        message_input = gr.Textbox(label="Message", placeholder="Type your message here...")
+        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)
+        chat_output = gr.Chatbot(label="Chatbot")
+        submit_button = gr.Button("Submit")
 
-    # Output components
-    chat_output = gr.Chatbot(label="Chatbot")
+        def update_components(model_key):
+            model_info = MODELS[model_key]
+            new_sizes = list(model_info["sizes"].keys())
+            return [
+                gr.Dropdown(choices=new_sizes, value=new_sizes[0]),
+                gr.Image(visible=model_info["is_vision"])
+            ]
 
-    # 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
+        model_key_selector.change(
+            fn=update_components,
+            inputs=model_key_selector,
+            outputs=[model_size_selector, image_input]
         )
-        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
-    )
+        submit_button.click(
+            fn=generate_response,
+            inputs=[
+                message_input,
+                image_input,
+                chat_output,
+                model_key_selector,
+                model_size_selector,
+                temperature_slider,
+                top_p_slider,
+                max_tokens_slider
+            ],
+            outputs=chat_output
+        )
 
-    # 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()
+create_interface().launch()