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


## update the model
tokenizer = AutoTokenizer.from_pretrained("cifope/nllb-fr-wol-wol-fr-3.3B")
model = AutoModelForSeq2SeqLM.from_pretrained("cifope/nllb-fr-wol-wol-fr-3.3B")

def translate(text):
  input_ids = tokenizer(text, return_tensors="pt").input_ids
  output_ids = model.generate(input_ids)
  output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
  return output_text

wolof2french = gr.Interface(fn=translate, 
             inputs="text", 
             outputs="text", 
             title="Traduction du wolof vers le français", 
             description="Ce modèle traduit le wolof vers le français. Entrez un texte en wolof et cliquez sur submit pour voir la traduction.", 
                          examples=[                       
                       ["Jàmm nga yendo !",  
                        
                        "Salut à toi !"], 
                       
                        ["Naka la mëna jënde crédit ci sama téléphone ?",  
                        
                        "Comment acheter un crédit sur mon téléphone ? "], 
                       
                       ["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."]
                       ],                 
             input_label="Texte en wolof", 
             output_label="Traduction en français")

wolof2french.launch() # share=True