import gradio as gr
from transformers import AutoTokenizer
from huggingface_hub import HfApi, login

api = HfApi()

# Define a function to calculate tokens
def count_tokens(llm_name, input_text, api_token):
    try:
        # Login using the API token if provided
        if api_token:
            login(api_token)
        
        # Load the tokenizer for the selected transformer-based model
        tokenizer = AutoTokenizer.from_pretrained(llm_name)
        tokens = tokenizer.encode(input_text)
        return f"Number of tokens: {len(tokens)}"
    except Exception as e:
        return f"Error: {str(e)}"

# Fetch model details including metadata (like tags)
models = list(api.list_models(task="text-generation"))

# Filter models that have the 'text-generation-inference' tag and 'text-generation' pipeline_tag
filtered_models = []
for model in models:
    model_info = api.model_info(model.modelId)
    if 'text-generation-inference' in model_info.tags and model_info.pipeline_tag == 'text-generation':
        filtered_models.append(model.modelId)

# Define custom CSS for a bluish theme and cursor pointer
custom_css = """
.gr-dropdown {
    cursor: pointer;
}
"""

# Set the default model to the first filtered model, or "gpt2" if there are no filtered models
default_model = filtered_models[0] if filtered_models else "gpt2"

# Create the Gradio interface
with gr.Blocks(css=custom_css) as demo:
    gr.HTML("<h1 style='text-align: center; color: #0078d7;'>Token Counter for Transformer-Based Models</h1>")
    gr.Markdown(
        "This app allows you to count the number of tokens in the input text "
        "using selected transformer-based models from Hugging Face."
    )
    with gr.Row():
        llm_dropdown = gr.Dropdown(choices=filtered_models, label="Select Transformer Model", value=default_model)
    with gr.Row():
        input_text = gr.Textbox(label="Enter your text")
        output = gr.Textbox(label="Token Count", interactive=False)
    with gr.Row():
        api_token_input = gr.Textbox(label="Enter Hugging Face API Token (if needed)", type="password", placeholder="Your API Token", interactive=True)
    with gr.Row():
        submit_btn = gr.Button("Calculate Tokens")

    submit_btn.click(count_tokens, inputs=[llm_dropdown, input_text, api_token_input], outputs=output)

# Launch the app
demo.launch(share=True, debug=True)