Spaces:
Runtime error
Runtime error
import gradio as gr | |
from datasets import load_dataset | |
imdb = load_dataset("imdb") | |
from transformers import AutoTokenizer, pipeline | |
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased") | |
def preprocess_function(examples): | |
return tokenizer(examples["text"], truncation=True) | |
tokenized_imdb = imdb.map(preprocess_function, batched=True) | |
from transformers import DataCollatorWithPadding | |
data_collator = DataCollatorWithPadding(tokenizer=tokenizer) | |
from transformers import AutoModelForSequenceClassification, TrainingArguments, Trainer | |
model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased", num_labels=2) | |
training_args = TrainingArguments( | |
output_dir="./results", | |
learning_rate=2e-5, | |
per_device_train_batch_size=16, | |
per_device_eval_batch_size=16, | |
num_train_epochs=0.01, | |
weight_decay=0.01, | |
) | |
trainer = Trainer( | |
model=model, | |
args=training_args, | |
train_dataset=tokenized_imdb["train"], | |
eval_dataset=tokenized_imdb["test"], | |
tokenizer=tokenizer, | |
data_collator=data_collator, | |
) | |
trainer.train() | |
def greet(text): | |
pipe = pipeline("sentiment-analysis", tokenizer=tokenizer, model=model) | |
return pipe(text)[0]['label'] | |
iface = gr.Interface(fn=greet, inputs=gr.inputs.Textbox(placeholder="Please enter the sentence...", lines=5), outputs="text") | |
iface.launch() |