import gradio as gr
from gradio.components import Dropdown, Textbox
from huggingface_hub import HfApi, ModelFilter
from transformers import pipeline

# Get the list of models from the Hugging Face Hub
api = HfApi()
models = api.list_models(author="jat-project", filter=ModelFilter(tags="text-generation"))
models_names = [model.modelId for model in models]

# Dictionary to store loaded models and their pipelines
model_pipelines = {}

# Load a default model initially
default_model_name = "jat-project/jat-small"


def generate_text(model_name, input_text):
    # Check if the selected model is already loaded
    if model_name not in model_pipelines:
        # Inform the user that the model is loading
        yield "Loading model..."

        # Load the model and create a pipeline if it's not already loaded
        generator = pipeline("text-generation", model=model_name, trust_remote_code=True)
        model_pipelines[model_name] = generator

    # Get the pipeline for the selected model
    generator = model_pipelines[model_name]

    # Inform the user that the text is being generated
    yield "Generating text..."

    # Generate text
    generated_text = generator(input_text, max_length=100)[0]["generated_text"]

    # Return the generated text
    yield generated_text


# Define the Gradio interface
iface = gr.Interface(
    fn=generate_text,  # Function to be called on user input
    inputs=[
        Dropdown(models_names, label="Select Model", value=default_model_name),  # Select model
        Textbox(lines=5, label="Input Text"),  # Textbox for entering text
    ],
    outputs=Textbox(label="Generated Text"),  # Textbox to display the generated text
    title="JAT Text Generation",  # Title of the interface
)

# Launch the Gradio interface
iface.launch(enable_queue=True)