from typing import List
import gradio as gr
from transformers import AutoModel, AutoTokenizer
from huggingface_hub import list_models

def get_collection_models(collection_name: str) -> List[str]:
    """Get a list of models from a specific Hugging Face collection."""
    models = list_models(author="de-Rodrigo",
                         filter=f"collections:{collection_name}")
    models_nofilter = list_models(author="de-Rodrigo")
    print("")
    print(models_nofilter)
    print("")
    model_names = [model.modelId for model in models]
    return model_names

def load_model(model_name: str):
    """Load a model from Hugging Face Hub."""
    model = AutoModel.from_pretrained(model_name)
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    return model, tokenizer

# Example processing function
def process_input(text: str, model_name: str) -> str:
    model, tokenizer = load_model(model_name)
    inputs = tokenizer(text, return_tensors="pt")
    outputs = model(**inputs)
    return f"Processed output with {model_name}"

# Create Gradio interface
def create_interface(collection_name: str):
    iface = gr.Interface(
        fn=process_input,
        inputs=[
            gr.Textbox(lines=2, placeholder="Enter some text", label="Input Text"),
            gr.Dropdown(choices=get_collection_models(collection_name), label="Select Model")
        ],
        outputs=gr.Textbox(label="Model Output"),
        title="Hugging Face Model Selector from Collection",
        description=f"Select a model from the '{collection_name}'.")
    return iface

# Specify the name of your collection
collection_name = "VrDU-Doctor"
iface = create_interface(collection_name)
iface.launch()