Añadiendo las traducciones

app.py

@@ -7,15 +7,69 @@ os.system('pip install --upgrade pip')
 os.system('pip install tensorflow')
 
 from transformers import pipeline
+from transformers import MarianMTModel, MarianTokenizer
+from nltk.tokenize import sent_tokenize
+from nltk.tokenize import LineTokenizer
+import math
+import torch
+import nltk
+nltk.download('punkt')
 
 docs = None
 
+if torch.cuda.is_available():  
+  dev = "cuda"
+else:  
+  dev = "cpu" 
+device = torch.device(dev)
 
 def request_pathname(files):
     if files is None:
         return [[]]
     return [[file.name, file.name.split('/')[-1]] for file in files]
 
+def traducir_parrafos(parrafos, tokenizer, model, tam_bloque=8, ):
+  parrafos_traducidos = []
+  for parrafo in parrafos:
+    frases = sent_tokenize(parrafo)
+    batches = math.ceil(len(frases) / tam_bloque)     
+    traducido = []
+    for i in range(batches):
+
+        bloque_enviado = frases[i*tam_bloque:(i+1)*tam_bloque]
+        model_inputs = tokenizer(bloque_enviado, return_tensors="pt", 
+                                 padding=True, truncation=True, 
+                                 max_length=500).to(device)
+        with torch.no_grad():
+            bloque_traducido = model.generate(**model_inputs)
+        traducido += bloque_traducido
+    traducido = [tokenizer.decode(t, skip_special_tokens=True) for t in traducido]
+    parrafos_traducidos += [" ".join(traducido)]
+  return parrafos_traducidos
+
+def traducir_es_en(texto):
+    mname = "Helsinki-NLP/opus-mt-es-en"
+    tokenizer = MarianTokenizer.from_pretrained(mname)
+    model = MarianMTModel.from_pretrained(mname)
+    model.to(device)
+    
+    lt = LineTokenizer()
+    batch_size = 8
+    parrafos = lt.tokenize(text_long)
+    par_tra = traducir_parrafos(parrafos, tokenizer, model) 
+    return "\n".join(par_tra)    
+
+def traducir_en_es(texto):
+    mname = "Helsinki-NLP/opus-mt-en-es"
+    tokenizer = MarianTokenizer.from_pretrained(mname)
+    model = MarianMTModel.from_pretrained(mname)
+    model.to(device)
+    
+    lt = LineTokenizer()
+    batch_size = 8
+    parrafos = lt.tokenize(text_long)
+    par_tra = traducir_parrafos(parrafos, tokenizer, model) 
+    return "\n".join(par_tra)
 
 def validate_dataset(dataset):
     global docs
@@ -27,19 +81,19 @@ def validate_dataset(dataset):
         return "⚠️Esperando documentos..."
 
 def do_ask(question, button, dataset):
-
     global docs
     docs_ready = dataset.iloc[-1, 0] != ""
     if button == "✨Listo✨" and docs_ready:
         for _, row in dataset.iterrows():
             path = row['filepath']
             text = Path(f'{path}').read_text()
+            text_en = traducir_es_en(text)
             question_answerer = pipeline("question-answering", model='distilbert-base-cased-distilled-squad')
             QA_input = {
-                'question': question,
-                'context': text
+                'question': traducir_es_en(question),
+                'context': text_en
             }
-            return question_answerer(QA_input)['answer']
+            return traducir_en_es(question_answerer(QA_input)['answer'])
     else:        
         return ""