Spaces:
Runtime error
Runtime error
File size: 2,360 Bytes
09be9ae 5114a54 09be9ae 5114a54 09be9ae b937959 7c322b8 09be9ae b937959 da6484f 09be9ae 249433b 5114a54 249433b 5114a54 249433b 5114a54 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
import torch
from transformers import RobertaTokenizer, RobertaForSequenceClassification, Trainer, TrainingArguments
from datasets import load_dataset, load_metric
import gradio as gr
# Carregar o dataset IMDB
dataset = load_dataset('imdb')
metric = load_metric('accuracy')
# 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 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
)
# Treinar o modelo
trainer.train()
# Avaliar o modelo
results = trainer.evaluate()
print(results)
# Salvar o modelo
model.save_pretrained('./model')
tokenizer.save_pretrained('./model')
# Função de inferência
def predict(text):
inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
outputs = model(**inputs)
predictions = torch.argmax(outputs.logits, dim=-1)
return "Positive" if predictions.item() == 1 else "Negative"
# Interface Gradio
iface = gr.Interface(
fn=predict,
inputs=gr.inputs.Textbox(lines=2, placeholder="Enter a movie review..."),
outputs="text",
title="IMDB Review Sentiment Analysis",
description="A simple Gradio interface to predict sentiment of IMDB movie reviews using a RoBERTa model."
)
if __name__ == "__main__":
iface.launch()
|