Update app.py

app.py

@@ -1,7 +1,12 @@
+# Importações necessárias
 import pandas as pd
 from sklearn.model_selection import train_test_split
-from transformers import AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments
+from torch.utils.data import Dataset
+import torch
+import gradio as gr
 
+# === PASSO 1: CARREGAR E PRÉ-PROCESSAR OS DADOS ===
 # Carregar os dados
 df = pd.read_csv("files/dados.csv")
 
@@ -25,4 +30,90 @@ tokenizer = AutoTokenizer.from_pretrained("neuralmind/bert-base-portuguese-cased
 train_encodings = tokenizer(train_texts, truncation=True, padding=True, max_length=128)
 test_encodings = tokenizer(test_texts, truncation=True, padding=True, max_length=128)
 
-# Restante do código para criar o dataset e fine-tuning...
+# Criar um dataset personalizado
+class CustomDataset(Dataset):
+    def __init__(self, encodings, labels):
+        self.encodings = encodings
+        self.labels = labels
+
+    def __len__(self):
+        return len(self.labels)
+
+    def __getitem__(self, idx):
+        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
+        item["labels"] = torch.tensor(self.labels[idx])
+        return item
+
+# Criar datasets
+train_dataset = CustomDataset(train_encodings, train_labels)
+test_dataset = CustomDataset(test_encodings, test_labels)
+
+# === PASSO 2: FINE-TUNING DO MODELO ===
+# Carregar o modelo pré-treinado para classificação
+model = AutoModelForSequenceClassification.from_pretrained(
+    "neuralmind/bert-base-portuguese-cased",
+    num_labels=3  # Número de classes (Baixa, Média, Alta)
+)
+
+# Configurar os argumentos de treinamento
+training_args = TrainingArguments(
+    output_dir="./results",
+    evaluation_strategy="epoch",
+    learning_rate=2e-5,
+    per_device_train_batch_size=8,
+    per_device_eval_batch_size=8,
+    num_train_epochs=3,
+    weight_decay=0.01,
+    logging_dir="./logs",
+    logging_steps=10,
+    save_strategy="epoch"
+)
+
+# Criar o Trainer
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=train_dataset,
+    eval_dataset=test_dataset
+)
+
+# Fine-tune o modelo
+print("Iniciando o fine-tuning do modelo...")
+trainer.train()
+print("Fine-tuning concluído!")
+
+# Salvar o modelo ajustado
+model.save_pretrained("./modelo-ajustado")
+tokenizer.save_pretrained("./modelo-ajustado")
+
+# === PASSO 3: INTEGRAR COM GRADIO ===
+# Carregar o modelo ajustado
+classifier = pipeline("text-classification", model="./modelo-ajustado")
+
+# Função para classificar a criticidade
+def classificar_criticidade(descricao):
+    resultado = classifier(descricao)[0]
+    label = resultado['label']
+    score = resultado['score']
+    
+    # Mapear os rótulos ajustados
+    if label == "LABEL_0":
+        return f"Criticidade: Baixa (Confiança: {score:.2f})"
+    elif label == "LABEL_1":
+        return f"Criticidade: Média (Confiança: {score:.2f})"
+    elif label == "LABEL_2":
+        return f"Criticidade: Alta (Confiança: {score:.2f})"
+    else:
+        return "Não foi possível determinar a criticidade."
+
+# Interface Gradio
+interface = gr.Interface(
+    fn=classificar_criticidade,
+    inputs=gr.Textbox(lines=2, placeholder="Descreva a compra..."),
+    outputs="text",
+    title="Classificador de Criticidade de Compra",
+    description="Insira a descrição da compra para receber uma classificação de criticidade."
+)
+
+# Iniciar a interface
+interface.launch()