Update app.py

app.py

@@ -6,12 +6,21 @@ client = InferenceClient("ulisesbravo/autotrain-nsuej-5ctie")
 
 def classify_text(text):
     # Faça a inferência no modelo de classificação
-    response = client(text)  # Utilize __call__ se o método predict não existir
+    response = client(text)  # Chama o cliente diretamente
     # Imprima a resposta para verificar sua estrutura
     print(response)  # Inspecione a resposta para entender a estrutura
+    
     # Extraia a classificação do response
-    # Ajuste isso conforme a estrutura do response
-    predicted_class = response['label'] if 'label' in response else "Classificação não encontrada"
+    # O modelo provavelmente retorna um array de logits, onde cada índice corresponde a uma classe
+    logits = response.get('logits')  # Ajuste isso conforme a estrutura do response
+    
+    if logits is not None:
+        # Obtém o índice da classe com a maior pontuação
+        predicted_index = logits.argmax().item()  # Converte para um índice
+        predicted_class = response['id2label'][str(predicted_index)]  # Acessa a classe usando id2label
+    else:
+        predicted_class = "Classificação não encontrada"
+
     return predicted_class
 
 # Interface Gradio