Update src/brain.py

src/brain.py

@@ -4,19 +4,18 @@ 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.")
+    inputs = tokenizer(query, return_tensors='pt', padding='max_length', truncation=True, max_length=512)
+    attention_mask = inputs['attention_mask']
+    input_ids = inputs['input_ids']
     
-    # Tokenização
-    inputs = tokenizer(query, return_tensors="pt", padding=True, truncation=True)
+    generated_ids = model.generate(
+        input_ids,
+        attention_mask=attention_mask,
+        max_length=len(input_ids[0]) + 100,  # Aumentar o limite de geração
+        temperature=0.7,  # Ajustar a criatividade
+        top_p=0.9,  # Usar nucleus sampling
+        no_repeat_ngram_size=2  # Evitar repetições desnecessárias
+    )
     
-    # Realizando a predição
-    outputs = model(**inputs)
-    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
+    generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+    return generated_text