Upload app.py

app.py

@@ -1,17 +1,15 @@
 import os
 import gradio as gr
-from huggingface_hub import list_repo_files, hf_hub_download
+from huggingface_hub import list_models, list_repo_files, hf_hub_download
 import subprocess
 
-# Constants
 HF_USER = "fbaldassarri"
 TEQ_KEYWORD = "TEQ"
 
 def list_teq_models():
-    # List all repos with TEQ in their name
-    from huggingface_hub import list_repos
-    repos = list_repos(HF_USER)
-    return [repo.id for repo in repos if TEQ_KEYWORD in repo.id]
+    # 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
@@ -48,24 +46,32 @@ def run_teq_inference(model_id, weights_file, config_file, base_model, prompt, m
         return output.split(marker)[-1].strip()
     return output
 
-# Gradio UI
-def ui():
+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 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")
-        model_id = gr.Dropdown(teq_models, label="Select TEQ Model")
-        weights_file = gr.Textbox(label="Weights File (.pt)")
-        config_file = gr.Textbox(label="Config File (.json)")
+        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)")
         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")
         debug = gr.Checkbox(label="Debug Mode")
         output = gr.Textbox(label="Generated Text", lines=10)
-        def update_files(model_id):
-            weights, configs = list_model_files(model_id)
-            return gr.update(choices=weights), gr.update(choices=configs)
-        model_id.change(update_files, inputs=model_id, outputs=[weights_file, config_file])
         run_btn = gr.Button("Run Inference")
+
+        model_id.change(
+            update_files,
+            inputs=model_id,
+            outputs=[weights_file, config_file]
+        )
         run_btn.click(
             run_teq_inference,
             inputs=[model_id, weights_file, config_file, base_model, prompt, max_new_tokens, debug],
@@ -74,4 +80,4 @@ def ui():
     return demo
 
 if __name__ == "__main__":
-    ui().launch()
+    build_ui().launch()