Upload app.py

app.py

@@ -17,6 +17,8 @@ def list_model_files(model_id):
     return weights, configs
 
 def run_teq_inference(model_id, weights_file, config_file, base_model, prompt, max_new_tokens, debug):
+    if not weights_file or not config_file:
+        return "Please select both a weights file and a config file."
     local_model_dir = f"./models/{model_id.replace('/', '_')}"
     os.makedirs(local_model_dir, exist_ok=True)
     hf_hub_download(model_id, weights_file, local_dir=local_model_dir)
@@ -42,9 +44,11 @@ def run_teq_inference(model_id, weights_file, config_file, base_model, prompt, m
 
 def update_files(model_id):
     weights, configs = list_model_files(model_id)
-    weights_val = weights[0] if weights else ""
-    configs_val = configs[0] if configs else ""
-    return weights, weights_val, configs, configs_val
+    # Set value to first item if available, else None
+    weights_val = weights[0] if weights else None
+    configs_val = configs[0] if configs else None
+    return gr.Dropdown(choices=weights, value=weights_val, label="Weights File (.pt)", interactive=True), \
+           gr.Dropdown(choices=configs, value=configs_val, label="Config File (.json)", interactive=True)
 
 def build_ui():
     teq_models = list_teq_models()
@@ -60,11 +64,11 @@ def build_ui():
         output = gr.Textbox(label="Generated Text", lines=10)
         run_btn = gr.Button("Run Inference")
 
-        # When model_id changes, update weights_file and config_file dropdowns
+        # Dynamically update the dropdowns for weights and config
         model_id.change(
             update_files,
             inputs=model_id,
-            outputs=[weights_file, weights_file, config_file, config_file]
+            outputs=[weights_file, config_file]
         )
         run_btn.click(
             run_teq_inference,