Update app.py

app.py

@@ -22,6 +22,22 @@ nltk.download('punkt')
 
 docs = None
 
+# Definimos los modelos:
+# Traducción
+mname = "Helsinki-NLP/opus-mt-es-en"
+tokenizer_es_en = MarianTokenizer.from_pretrained(mname)
+model_es_en = MarianMTModel.from_pretrained(mname)
+model_es_en.to(device)
+
+mname = "Helsinki-NLP/opus-mt-en-es"
+tokenizer_en_es = MarianTokenizer.from_pretrained(mname)
+model_en_es = MarianMTModel.from_pretrained(mname)
+model_en_es.to(device)
+
+lt = LineTokenizer()
+
+# Responder preguntas
+
 tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/multi-qa-mpnet-base-dot-v1")
 model = AutoModel.from_pretrained("sentence-transformers/multi-qa-mpnet-base-dot-v1").to(device).eval()
 tokenizer_ans = AutoTokenizer.from_pretrained("deepset/roberta-large-squad2")
@@ -41,6 +57,35 @@ def validate_dataset(dataset):
     else:
         return "⚠️Esperando documentos..."
 
+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):
+    parrafos = lt.tokenize(texto)
+    par_tra = traducir_parrafos(parrafos, tokenizer_es_en, model_es_en) 
+    return "\n".join(par_tra)    
+
+def traducir_en_es(texto):
+    parrafos = lt.tokenize(texto)
+    par_tra = traducir_parrafos(parrafos, tokenizer_en_es, model_en_es) 
+    return "\n".join(par_tra)
+
 def request_pathname(files):
     if files is None:
         return [[]]
@@ -101,8 +146,9 @@ def encode_docs(docs,maxlen = 64, stride = 32):
     return embeddings, spans, file_names
    
 def predict(query,data):
+    query = traducir_es_en(query)
     name_to_save = data.name.split("/")[-1].split(".")[0][:-8]
-    k=20
+    k=2
     st = str([query,name_to_save])
     st_hashed = str(hashlib.sha256(st.encode()).hexdigest()) #just to speed up examples load
     hist = st + " " + st_hashed 
@@ -114,7 +160,9 @@ def predict(query,data):
         list_outputs = []
         for i in range(k):
             temp = [df.iloc[n] for n in range(k)][i]
-            tupla = (temp.Respuesta, temp.Contexto, temp.Probabilidades)
+            tupla = (traducir_en_es(temp.Respuesta), 
+                     traducir_en_es(temp.Contexto), 
+                     traducir_en_es(temp.Probabilidades))
             # text = ''
             # text += 'Probabilidades: '+ temp.Probabilidades + '\n\n' 
             # text += 'Respuesta: ' +temp.Respuesta + '\n\n' 
@@ -139,6 +187,8 @@ def predict(query,data):
         text = text.replace("\r", " ")
         text = text.replace("\n", " ")
         text = text.replace(" . "," ")
+
+        text = traducir_es_en(text)
         
         doc_emb, doc_text, file_names = encode_docs((name_to_save,text),maxlen = 64, stride = 32)
         
@@ -194,7 +244,9 @@ def predict(query,data):
     list_outputs = []
     for i in range(k):
         temp = [df.iloc[n] for n in range(k)][i]
-        tupla = (temp.Respuesta, temp.Contexto, temp.Probabilidades)
+        tupla = (traducir_en_es(temp.Respuesta), 
+                 traducir_en_es(temp.Contexto), 
+                 traducir_en_es(temp.Probabilidades))
         # text = ''
         # text += 'Probabilidades: '+ temp.Probabilidades + '\n\n' 
         # text += 'Respuesta: ' +temp.Respuesta + '\n\n' 
@@ -248,7 +300,7 @@ with gr.Blocks() as demo:
     ask.click(fn=predict, inputs=[query,
                                  file], outputs=[answer, context, prob])
 
-    examples = ["¿Cuándo suelen comenzar las adicciones?","Entrevista Miguel Ruiz.txt"]
+    gr.Interface.load(examples = ["¿Cuándo suelen comenzar las adicciones?","Entrevista Miguel Ruiz.txt"])
 
 demo.queue(concurrency_count=20)
 demo.launch(show_error=True)