from gradio import Interface
import gradio as gr
import aranizer
from aranizer import aranizer_bpe50k, aranizer_bpe64k, aranizer_bpe86k, aranizer_sp32k, aranizer_sp50k, aranizer_sp64k, aranizer_sp86k
from transformers import AutoTokenizer, logging
from huggingface_hub import HfApi, login
import os

# Retrieve your Hugging Face token from the environment variable
HF_TOKEN = os.getenv('HF_TOKEN')

if HF_TOKEN:
    api = HfApi()
    api.set_access_token(HF_TOKEN)
    login(token=HF_TOKEN)

# Load additional tokenizers from transformers
gpt_13b_tokenizer = AutoTokenizer.from_pretrained("FreedomIntelligence/AceGPT-13B")
gpt_7b_tokenizer = AutoTokenizer.from_pretrained("FreedomIntelligence/AceGPT-7B")
jais_13b_tokenizer = AutoTokenizer.from_pretrained("inception-mbzuai/jais-13b")
arabert_tokenizer = AutoTokenizer.from_pretrained("aubmindlab/bert-base-arabertv2")

# Try to load the gated tokenizer
try:
    meta_llama_tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B")
except Exception as e:
    meta_llama_tokenizer = None
    logging.warning(f"Could not load meta-llama/Meta-Llama-3-8B tokenizer: {e}")

cohere_command_r_v01_tokenizer = AutoTokenizer.from_pretrained("CohereForAI/c4ai-command-r-v01")
cohere_command_r_plus_tokenizer = AutoTokenizer.from_pretrained("CohereForAI/c4ai-command-r-plus")

# List of available tokenizers and a dictionary to load them
tokenizer_options = [
    "aranizer_bpe50k", "aranizer_bpe64k", "aranizer_bpe86k",
    "aranizer_sp32k", "aranizer_sp50k", "aranizer_sp64k", "aranizer_sp86k",
    "FreedomIntelligence/AceGPT-13B",  
    "FreedomIntelligence/AceGPT-7B",   
    "inception-mbzuai/jais-13b",
    "aubmindlab/bert-base-arabertv2",
    "CohereForAI/c4ai-command-r-v01",
    "CohereForAI/c4ai-command-r-plus"   
]

if meta_llama_tokenizer:
    tokenizer_options.append("meta-llama/Meta-Llama-3-8B")

tokenizers = {
    "aranizer_bpe50k": aranizer_bpe50k.get_tokenizer,
    "aranizer_bpe64k": aranizer_bpe64k.get_tokenizer,
    "aranizer_bpe86k": aranizer_bpe86k.get_tokenizer,
    "aranizer_sp32k": aranizer_sp32k.get_tokenizer,
    "aranizer_sp50k": aranizer_sp50k.get_tokenizer,
    "aranizer_sp64k": aranizer_sp64k.get_tokenizer,
    "aranizer_sp86k": aranizer_sp86k.get_tokenizer,
    "FreedomIntelligence/AceGPT-13B": lambda: gpt_13b_tokenizer,
    "FreedomIntelligence/AceGPT-7B": lambda: gpt_7b_tokenizer,
    "inception-mbzuai/jais-13b": lambda: jais_13b_tokenizer,
    "aubmindlab/bert-base-arabertv2": lambda: arabert_tokenizer,
    "CohereForAI/c4ai-command-r-v01": lambda: cohere_command_r_v01_tokenizer,
    "CohereForAI/c4ai-command-r-plus": lambda: cohere_command_r_plus_tokenizer
}

if meta_llama_tokenizer:
    tokenizers["meta-llama/Meta-Llama-3-8B"] = lambda: meta_llama_tokenizer

def compare_tokenizers(tokenizer_name, text):
    # Handle the transformer tokenizers separately due to API differences
    if tokenizer_name in [
        "FreedomIntelligence/AceGPT-13B", "FreedomIntelligence/AceGPT-7B", 
        "inception-mbzuai/jais-13b", "aubmindlab/bert-base-arabertv2",
        "meta-llama/Meta-Llama-3-8B", "CohereForAI/c4ai-command-r-v01", "CohereForAI/c4ai-command-r-plus"
    ]:
        tokenizer = tokenizers[tokenizer_name]()
        tokens = tokenizer.tokenize(text)
        encoded_output = tokenizer.encode(text, add_special_tokens=True)
        decoded_text = tokenizer.decode(encoded_output, skip_special_tokens=True)
    else:
        # AraNizer tokenizers
        tokenizer = tokenizers[tokenizer_name]()
        tokens = tokenizer.tokenize(text)
        encoded_output = tokenizer.encode(text, add_special_tokens=True)
        decoded_text = tokenizer.decode(encoded_output)
    
    # Prepare the results to be displayed in HTML format
    tokens_arabic = [token.encode('utf-8').decode('utf-8') if isinstance(token, bytes) else token for token in tokens]
    results_html = f"""
    <div>
        <h3>Tokenizer: {tokenizer_name}</h3>
        <p><strong>Tokens:</strong> {tokens_arabic}</p>
        <p><strong>Encoded:</strong> {encoded_output}</p>
        <p><strong>Decoded:</strong> {decoded_text}</p>
    </div>
    """
    return results_html

# Define the Gradio interface components with a dropdown for model selection
inputs_component = [
    gr.Dropdown(choices=tokenizer_options, label="Select Tokenizer"),
    gr.Textbox(lines=2, placeholder="اكتب النص هنا...", label="Input Text")
]

outputs_component = gr.HTML(label="Results")

# Setting up the interface
iface = Interface(
    fn=compare_tokenizers, 
    inputs=inputs_component, 
    outputs=outputs_component, 
    title="Arabic Tokenizer Arena",
    live=True
)

# Launching the Gradio app
iface.launch()