# Import necessary libraries
import gradio as gr
import torch
import spaces
from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM, pipeline


# Load Goldfish model for Dhivehi
model_name = 'div_thaa_full'
HF_CACHE = '.hf_cache'

# Load model
goldfish_model = 'goldfish-models/' + model_name
config = AutoConfig.from_pretrained(goldfish_model, cache_dir=HF_CACHE)
tokenizer = AutoTokenizer.from_pretrained(goldfish_model, cache_dir=HF_CACHE)
model = AutoModelForCausalLM.from_pretrained(goldfish_model, config=config, cache_dir=HF_CACHE)

if torch.cuda.is_available():
    model = model.cuda()  # Load onto GPU


# Create text generation pipeline
text_generator = pipeline('text-generation', model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1)

# Function to generate text
@spaces.GPU
def generate_text(input_text):
    output = text_generator(input_text, max_new_tokens=25, add_special_tokens=False, do_sample=True)
    return output[0]['generated_text']

# Create Gradio interface

with gr.Blocks(css="""
    .thaana textarea {
        font-size: 18px !important;
        font-family: 'MV_Faseyha', 'Faruma', 'A_Faruma', 'Noto Sans Thaana', 'MV Boli';
        line-height: 1.8 !important;
    }
""") as demo:
    gr.Markdown("# Demo Dhivehi Text Generator")
    gr.Markdown("Generate text in Dhivehi language. This model is trained to generate coherent text based on the input prompt.")
    
    with gr.Row():
        input_text = gr.Textbox(
            lines=2,
            label="Enter Dhivehi Text",
            rtl=True,
            elem_classes="thaana"
        )
        output_text = gr.Textbox(
            lines=2,
            rtl=True,
            elem_classes="thaana"
        )

    generate_btn = gr.Button("Generate")
    generate_btn.click(
        fn=generate_text,
        inputs=input_text,
        outputs=output_text
    )

    gr.Markdown("""
        Model: Goldfish is a suite of monolingual language models trained for 350 languages. 
        This model is the Dhivehi (Thaana script). For more details, visit the 
        [Goldfish Models GitHub repository](https://github.com/tylerachang/goldfish).
    """)

    examples = gr.Examples(
        examples=[
            ["ދިވެހިރާއްޖެ"],
            ["އެމެރިކާ އިންތިޚާބު"],
            ["ސަލާމް"],
            ["ދުނިޔޭގެ ސިއްޙަތު ޖަމްޢިއްޔާ"],
            ["ޤަދީމީ ސަގާފަތް"],
            ["ޑިމޮކްރަސީ"]
        ],
        inputs=input_text,
        outputs=output_text,
        fn=generate_text
    )

if __name__ == "__main__":
    demo.launch()