import gradio as gr
import spaces
import torch
import os
from transformers import AutoTokenizer, AutoModelForCausalLM
from datetime import datetime

model_id = "BSC-LT/SalamandraTA-7B-instructed-Aranese"
token = os.getenv("AudreyVM")

# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(model_id, token=token)

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype=torch.bfloat16, 
    token=token
)

languages = sorted([ 'Aragonese', 'Asturian', 'Basque', 'Bulgarian', 'Catalan', 'Catalan_Valencian', 'Croatian', 'Czech', 'Danish', 'Dutch', 'English', 'Estonian',
             'Finnish', 'French', 'Galician', 'German', 'Greek', 'Hungarian', 'Irish', 'Italian', 'Latvian', 'Lithuanian', 'Maltese', 'Norwegian Bokmål',
             'Norwegian Nynorsk', 'Occitan', 'Aranese', 'Polish', 'Portuguese', 'Romanian', 'Russian', 'Serbian_Cyrillic', 'Slovak', 'Slovenian', 'Spanish', 'Swedish',
             'Ukrainian', 'Welsh' ])

@spaces.GPU(duration=120)
def generate_output(source, target, input_text, mt_text=None):
    date_string = datetime.today().strftime('%Y-%m-%d')


    sentences = input_text.split('\n')
    #sentences = [s for s in sentences if len(s.strip()) > 0]
    generated_text = []

    for sentence in sentences:
        sentence = sentence.strip()

        if len(sentence) == 0:
            # Preserve empty lines
            generated_text.append('')
            continue
        
        prompt = f"Translate the following text from {source} into {target}.\n{source}: {sentence.strip()} \n{target}:"
    
        messages = [{"role": "user", "content": prompt}]
        final_prompt = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
            date_string=date_string
        )
    
        inputs = tokenizer(final_prompt, return_tensors="pt", add_special_tokens=False).to(model.device)
        input_length = inputs.input_ids.shape[1]
    
        output = model.generate(
            input_ids=inputs.input_ids,
            max_new_tokens=4000,
            early_stopping=True,
            num_beams=1
        )
    
        decoded = tokenizer.decode(output[0, input_length:], skip_special_tokens=True).strip()
        generated_text.append(decoded)
        
    return '\n'.join(generated_text), ""

with gr.Blocks() as demo:
    gr.Markdown("# 🦎 SalamandraTA 7B - Aranese Demo")


    with gr.Row():
        source_lang = gr.Dropdown(choices=languages, value="Catalan", label="Source Language")
        target_lang = gr.Dropdown(choices=languages, value="Aranese", label="Target Language")

        if source_lang == 'Catalan_Valencian':
            source_lang = 'Valencian'

        if target_lang == 'Catalan_Valencian':
            target_lang = 'Valencian'

    input_textbox = gr.Textbox(lines=6, placeholder="Enter source text here", label="Input Text")
    output_textbox = gr.Textbox(lines=6, label="Output")

    info_label = gr.HTML("")
    translate_btn = gr.Button("Translate")
    translate_btn.click(generate_output, inputs=[source_lang, target_lang, input_textbox], outputs=[output_textbox, info_label])

    gr.Examples(
        examples=[
            ["Catalan", "Aranese", "Als antics egipcis del període de l'Imperi Nou els fascinaven els monuments dels seus predecessors, que llavors tenien més de mil anys.", ""],
            
        ],
        inputs=[source_lang, target_lang, input_textbox]
    )

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