Upload app.py

app.py

@@ -1,29 +1,18 @@
+# Instalar as dependências manualmente, se necessário
+import os
+os.system("pip install simpletransformers")
+
 import gradio as gr
-from transformers import BertTokenizer, BertForQuestionAnswering
-import torch
+from simpletransformers.question_answering import QuestionAnsweringModel
 
-# Carregando o modelo a partir do checkpoint e o tokenizer
-tokenizer = BertTokenizer.from_pretrained("outputs/bert/best_model")
-model = BertForQuestionAnswering.from_pretrained("outputs/bert/best_model")
+# Load model from training checkpoint
+model = QuestionAnsweringModel("bert", "outputs/bert/best_model", use_cuda=False)
 
 # Função para realizar predições
 def perform_prediction(context, question):
-    inputs = tokenizer.encode_plus(question, context, return_tensors="pt")
-    input_ids = inputs["input_ids"]
-    attention_mask = inputs["attention_mask"]
-
-    with torch.no_grad():
-        outputs = model(input_ids=input_ids, attention_mask=attention_mask)
-        start_scores, end_scores = outputs.start_logits, outputs.end_logits
-    
-    # Encontrando a posição da resposta
-    answer_start = torch.argmax(start_scores)
-    answer_end = torch.argmax(end_scores) + 1
-    
-    # Decodificando a resposta
-    answer = tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens(input_ids[0][answer_start:answer_end]))
-    
-    return answer
+    to_predict = [{"context": context, "qas": [{"question": question, "id": "0"}]}]
+    answers, probabilities = model.predict(to_predict, n_best_size=2)
+    return answers[0]["answer"]
 
 # Criando a interface do Gradio
 iface = gr.Interface(
@@ -38,4 +27,4 @@ iface = gr.Interface(
 )
 
 # Iniciando a interface
-iface.launch()
+iface.launch()