import gradio as gr
import torch
import os
from huggingface_hub import login
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCausalLM

login(token="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")
os.environ["HF_TOKEN"] = "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"

tokenizer = AutoTokenizer.from_pretrained("papasega/m2m100_418M_finetune_VF_dataV5_100_Epochs_wo-to-fr")
model = AutoModelForSeq2SeqLM.from_pretrained("papasega/m2m100_418M_finetune_VF_dataV5_100_Epochs_wo-to-fr")

# gpt2_model = AutoModelForCausalLM.from_pretrained("gpt2")
# gpt2_tokenizer = AutoTokenizer.from_pretrained("gpt2")

gpt2_tokenizer = AutoTokenizer.from_pretrained("bigscience/bloom-560m")
gpt2_model = AutoModelForCausalLM.from_pretrained("bigscience/bloom-560m")

# gpt2_tokenizer =  AutoTokenizer.from_pretrained("bofenghuang/vigogne-2-7b-instruct")
# gpt2_model = AutoModelForCausalLM.from_pretrained("bofenghuang/vigogne-2-7b-instruct")

def generate_text(text):
    input_ids = tokenizer(text, return_tensors="pt").input_ids
    output_ids = model.generate(input_ids)
    translation = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    gpt2_input = gpt2_tokenizer(translation, return_tensors="pt")
    gpt2_output = gpt2_model.generate(**gpt2_input)
    generated_text = gpt2_tokenizer.batch_decode(gpt2_output, max_length=1024, skip_special_tokens=True)[0]

    return translation, generated_text


wolof2french = gr.Interface(fn=generate_text, 
             inputs="text", 
             outputs=["text", "text"], 
             title="Génération du texte en français à partir du wolof", 
             description="Ce modèle traduit le wolof vers le français, puis génère du texte en français à partir de la traduction. Entrez un texte en wolof et cliquez sur submit pour voir la traduction et le texte généré.", 
                          examples=[                       
                       ["Jàmm nga yendo !",  
                        
                        "Bonne journée ! ", 
                        
                        "Bonne journée ! Passe le bonjour à ton."], 
                       
                        ["Naka la mëna jënde crédit ci sama téléphone ?",  
                        
                        "Comment acheter un crédit sur mon téléphone ? ", 
                        
                        "Comment acheter un crédit sur mon téléphone ? Il faut aller dans une boutique ou utiliser une application."], 
                       
                       ["Ngir xool ki ci lóru da nga wara njëkka saytu barab bi ngir àar sa kaarànge.", 
                                     
                        "Avant d'examiner la victime, assurez-votre sécurité en inspectant les lieux.", 
                        
                        "Avant d'examiner la victime, assurez-votre sécurité en inspectant les lieux. Ne touchez pas à des objets suspects ou dangereux. Appelez les secours si nécessaire."]
                       ],                 
             input_label="Texte en wolof", 
             output_label=["Traduction en français", "Texte généré en français"])

wolof2french.launch()
#wolof2french.close()