import gradio as gr
from random import random
from aranizer import aranizer_bpe32k, aranizer_bpe50k, aranizer_bpe64k, aranizer_bpe86k, aranizer_sp32k, aranizer_sp50k, aranizer_sp64k, aranizer_sp86k

def load_tokenizer(tokenizer_choice):
    # Dictionary mapping tokenizer choice to actual tokenizer initializer
    tokenizer_map = {
        "aranizer_bpe32k": aranizer_bpe32k.get_tokenizer(),
        "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(),
    }

    return tokenizer_map.get(tokenizer_choice, None)

def tokenize_and_encode_and_embed(text, tokenizer_choice):
    tokenizer = load_tokenizer(tokenizer_choice)  # Use load_tokenizer here
    if tokenizer:
        # Example methods. Replace with actual methods from your tokenizer
        tokens = tokenizer.tokenize(text)
        encoded_output = tokenizer.encode(text, add_special_tokens=True)
        decoded_text = tokenizer.decode(encoded_output)

        # Example embedding (replace with actual embedding generation from your model)
        embeddings = [random() for _ in range(10)]  # Example 10-dimensional embedding vector

        return " ".join(tokens), str(encoded_output), decoded_text, embeddings
    else:
        return "Tokenizer not loaded correctly", "", "", []

demo = gr.Interface(
    fn=tokenize_and_encode_and_embed,
    inputs=[
        gr.Textbox(lines=5, label="النص العربي"),
        gr.Dropdown(choices=["aranizer_bpe32k", "aranizer_bpe50k", "aranizer_bpe64k", "aranizer_bpe86k", "aranizer_sp32k", "aranizer_sp50k", "aranizer_sp64k", "aranizer_sp86k"], label="اختر المحلل اللفظي")
    ],
    outputs=[
        gr.Text(label="Tokens"),
        gr.Text(label="Encoded Output"),
        gr.Text(label="Decoded Text"),
        gr.Text(label="Embeddings (Example Vector)")
    ],
    title="مقارنة المحللات اللفظية وعمليات التضمين للنص العربي",
    description="حدد نوع المحلل اللفظي وأدخل نصًا لرؤية النتائج ومتجه التضمين.",
)

demo.launch()