Update src/brain.py

src/brain.py

@@ -1,11 +1,12 @@
-from transformers import AutoTokenizer, AutoModel
+from transformers import BertTokenizer, BertForSequenceClassification
 
-tokenizer = AutoTokenizer.from_pretrained('juridics/bertimbaulaw-base-portuguese-sts-scale')
-model = AutoModel.from_pretrained('juridics/bertimbaulaw-base-portuguese-sts-scale')
+tokenizer = BertTokenizer.from_pretrained('juridics/bertimbaulaw-base-portuguese-sts-scale')
+model = BertForSequenceClassification.from_pretrained('juridics/bertimbaulaw-base-portuguese-sts-scale')
 
-def generate_answers(query):
-    input_ids = tokenizer.encode(query, return_tensors='pt')
-    max_length = input_ids.shape[1] + 100
-    generated_ids = model.generate(input_ids, max_length=max_length)
-    generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
-    return generated_text
+def generate_answers(query, df):
+    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
+    predicted_label = labels[prediction]
+    return predicted_label