Update src/brain.py

src/brain.py

@@ -3,10 +3,10 @@ from transformers import BertTokenizer, BertForSequenceClassification
 tokenizer = BertTokenizer.from_pretrained('juridics/bertimbaulaw-base-portuguese-sts-scale')
 model = BertForSequenceClassification.from_pretrained('juridics/bertimbaulaw-base-portuguese-sts-scale')
 
-def generate_answers(query, df):
+def generate_answers(query, data):
     inputs = tokenizer(query, return_tensors="pt", padding=True, truncation=True)
     outputs = model(**inputs)
     prediction = torch.argmax(outputs.logits, dim=1)
-    labels = df.columns  # Substitua com suas etiquetas reais
+    labels = data.columns  # Substitua com suas etiquetas reais
     predicted_label = labels[prediction]
     return predicted_label