Update app.py

app.py

@@ -1,10 +1,10 @@
 import pandas as pd
 import gradio as gr
-from transformers import AutoTokenizer, AutoModel
+from transformers import BertTokenizer, BertForSequenceClassification
 import torch
 
-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')
 
 # Dados iniciais
 data = {
@@ -32,11 +32,12 @@ def get_gpt_response(query):
     {csv_data}
     
     """
-    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
+    inputs = tokenizer(query, return_tensors="pt", padding=True, truncation=True)
+    outputs = model(**inputs)
+    prediction = torch.argmax(outputs.logits, dim=1)
+    labels = ["Label1", "Label2"]  # Substitua com suas etiquetas reais
+    predicted_label = labels[prediction]
+    return predicted_label
 
 
 def ask_question(pergunta):