Update app.py

app.py

@@ -1,26 +1,58 @@
-import gradio as gr
-from transformers import pipeline
-
-# Carregar o modelo pré-treinado diretamente do Hugging Face
-ner_model = pipeline(
-    "ner", 
-    model="bert-base-cased",
-    tokenizer="bert-base-cased"
+import torch
+from transformers import RobertaTokenizer, RobertaForSequenceClassification, Trainer, TrainingArguments
+from datasets import load_dataset, load_metric
+
+# Carregar o dataset IMDB
+dataset = load_dataset('imdb')
+
+# Carregar o tokenizer e o modelo RoBERTa
+tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
+model = RobertaForSequenceClassification.from_pretrained('roberta-base')
+
+# Tokenizar os dados
+def preprocess_function(examples):
+    return tokenizer(examples['text'], padding='max_length', truncation=True)
+
+tokenized_datasets = dataset.map(preprocess_function, batched=True)
+
+# Preparar o data collator
+from transformers import DataCollatorWithPadding
+data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
+
+# Configurar os argumentos de treinamento
+training_args = TrainingArguments(
+    output_dir='./results',
+    evaluation_strategy='epoch',
+    per_device_train_batch_size=8,
+    per_device_eval_batch_size=8,
+    num_train_epochs=3,
+    weight_decay=0.01,
 )
 
-# Definir a função para a interface do Gradio
-def analyze_ner(text):
-    result = ner_model(text)
-    return result
-
-# Criar a interface do Gradio
-iface = gr.Interface(
-    fn=analyze_ner, 
-    inputs="text", 
-    outputs="json",
-    title="Reconhecimento de Entidades Nomeadas",
-    description="Digite um texto para identificar entidades nomeadas."
+# Definir a função de métricas
+def compute_metrics(eval_pred):
+    logits, labels = eval_pred
+    predictions = torch.argmax(logits, dim=-1)
+    return metric.compute(predictions=predictions, references=labels)
+
+# Definir o Trainer
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=tokenized_datasets['train'],
+    eval_dataset=tokenized_datasets['test'],
+    tokenizer=tokenizer,
+    data_collator=data_collator,
+    compute_metrics=compute_metrics
 )
 
-# Lançar a interface do Gradio com um link público
-iface.launch(share=True)
+# Treinar o modelo
+trainer.train()
+
+# Avaliar o modelo
+results = trainer.evaluate()
+print(results)
+
+# Salvar o modelo
+model.save_pretrained('./model')
+tokenizer.save_pretrained('./model')