Upload app.py

app.py

@@ -7,24 +7,20 @@ HF_USER = "fbaldassarri"
 TEQ_KEYWORD = "TEQ"
 
 def list_teq_models():
-    # List all models for the user, filter those with "TEQ" in the name
     models = list_models(author=HF_USER)
     return [model.modelId for model in models if TEQ_KEYWORD in model.modelId]
 
 def list_model_files(model_id):
-    # List files in the repo that are likely to be weights/config
     files = list_repo_files(model_id)
     weights = [f for f in files if f.endswith('.pt')]
     configs = [f for f in files if f.endswith('.json')]
     return weights, configs
 
 def run_teq_inference(model_id, weights_file, config_file, base_model, prompt, max_new_tokens, debug):
-    # Download files if not present
     local_model_dir = f"./models/{model_id.replace('/', '_')}"
     os.makedirs(local_model_dir, exist_ok=True)
-    weights_path = hf_hub_download(model_id, weights_file, local_dir=local_model_dir)
-    config_path = hf_hub_download(model_id, config_file, local_dir=local_model_dir)
-    # Call teq_inference.py as a subprocess for isolation
+    hf_hub_download(model_id, weights_file, local_dir=local_model_dir)
+    hf_hub_download(model_id, config_file, local_dir=local_model_dir)
     cmd = [
         "python", "teq_inference.py",
         "--model_dir", local_model_dir,
@@ -38,9 +34,7 @@ def run_teq_inference(model_id, weights_file, config_file, base_model, prompt, m
     if debug:
         cmd.append("--debug")
     result = subprocess.run(cmd, capture_output=True, text=True)
-    # Extract generated text from logs
     output = result.stdout + "\n" + result.stderr
-    # Try to find the generated text in logs
     marker = "Generated text:"
     if marker in output:
         return output.split(marker)[-1].strip()
@@ -48,18 +42,21 @@ 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)
+    # Default to first file if available, else empty string
     weights_val = weights[0] if weights else ""
     configs_val = configs[0] if configs else ""
-    return gr.Dropdown.update(choices=weights, value=weights_val), gr.Dropdown.update(choices=configs, value=configs_val)
+    return (
+        gr.Dropdown.update(choices=weights, value=weights_val),
+        gr.Dropdown.update(choices=configs, value=configs_val)
+    )
 
 def build_ui():
     teq_models = list_teq_models()
     with gr.Blocks() as demo:
         gr.Markdown("# TEQ Quantized Model Inference Demo")
-        with gr.Row():
-            model_id = gr.Dropdown(teq_models, label="Select TEQ Model")
-            weights_file = gr.Dropdown(choices=[], label="Weights File (.pt)")
-            config_file = gr.Dropdown(choices=[], label="Config File (.json)")
+        model_id = gr.Dropdown(teq_models, label="Select TEQ Model", interactive=True)
+        weights_file = gr.Dropdown(choices=[], label="Weights File (.pt)", interactive=True)
+        config_file = gr.Dropdown(choices=[], label="Config File (.json)", interactive=True)
         base_model = gr.Textbox(label="Base Model Name", value="facebook/opt-350m")
         prompt = gr.Textbox(label="Prompt", value="Once upon a time, a little girl")
         max_new_tokens = gr.Slider(10, 512, value=100, label="Max New Tokens")
@@ -67,6 +64,7 @@ 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
         model_id.change(
             update_files,
             inputs=model_id,