from transformers import AutoTokenizer
import gradio as gr
import os

# Retrieve the Hugging Face token from secrets
huggingface_token = os.getenv("HUGGINGFACE_TOKEN")

def tokenize(input_text):
    palmyra_x_003_tokens = len(palmyra_x_003_tokenizer(input_text, add_special_tokens=True)["input_ids"])
    gpt2_tokens = len(gpt2_tokenizer(input_text, add_special_tokens=True)["input_ids"])
    palmyra_x_004_tokens = len(palmyra_x_004_tokenizer(input_text, add_special_tokens=True)["input_ids"])

    results = {
        "Palmyra-X-004": palmyra_x_004_tokens,
        "Palmyra-Fin & Med": palmyra_x_003_tokens,
        "Palmyra-X-003": gpt2_tokens
    }

    # Sort the results in descending order based on token length
    sorted_results = sorted(results.items(), key=lambda x: x[1], reverse=True)

    return "\n".join([f"{model}: {tokens}" for model, tokens in sorted_results])


if __name__ == "__main__":
    palmyra_x_003_tokenizer = AutoTokenizer.from_pretrained("wassemgtk/palmyra-x-003-tokenizer", token=huggingface_token)
    gpt2_tokenizer = AutoTokenizer.from_pretrained("gpt2")
    palmyra_x_004_tokenizer = AutoTokenizer.from_pretrained("wassemgtk/palmyra-x-004-tokenizer", token=huggingface_token)

    iface = gr.Interface(fn=tokenize, inputs=gr.Textbox(label="Input Text", lines=19), outputs="text")
    iface.launch()