Update src/brain.py

src/brain.py

@@ -4,9 +4,19 @@ tokenizer = BertTokenizer.from_pretrained('juridics/bertimbaulaw-base-portuguese
 model = BertForSequenceClassification.from_pretrained('juridics/bertimbaulaw-base-portuguese-sts-scale')
 
 def generate_answers(query):
+    # Garantindo que a query é uma string
+    if not isinstance(query, str):
+        raise ValueError("A entrada para a função generate_answers deve ser uma string.")
+    
+    # Tokenização
     inputs = tokenizer(query, return_tensors="pt", padding=True, truncation=True)
+    
+    # Realizando a predição
     outputs = model(**inputs)
-    prediction = torch.argmax(outputs.logits, dim=1)
-    labels = ['ds','real','Group']
-    predicted_label = labels[prediction]
+    prediction = torch.argmax(outputs.logits, dim=1).item()  # Converter tensor para um inteiro
+    
+    # Labels devem corresponder ao número de classes do modelo
+    labels = ['ds', 'real', 'Group']
+    predicted_label = labels[prediction]  # Usando o índice para acessar a label
+    
     return predicted_label