File size: 1,573 Bytes
8225a4d
 
4d59358
b2f107f
730ad9d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8225a4d
730ad9d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4d59358
 
786fe91
 
 
 
 
 
 
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
# -*- coding: utf-8 -*-

import gradio as gr
from transformers import pipeline
from transformers import AutoTokenizer
from datasets import load_dataset
from transformers import DataCollatorWithPadding

raw_datasets = load_dataset("glue", "sst2")
raw_datasets
checkpoint = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
def tokenize_function(example):
  return tokenizer(example["sentence"], truncation=True)

tokenized_datasets = raw_datasets.map(tokenize_function, batched=True,remove_columns=['idx','sentence'])
tokenized_datasets


data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

from transformers import TrainingArguments
from transformers import AutoModelForSequenceClassification
from datasets import load_metric
from transformers import Trainer
import numpy as np

training_args = TrainingArguments("test-trainer", evaluation_strategy="epoch")
model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)

def compute_metrics(eval_preds):
    metric = load_metric("glue", "sst2")
    logits, labels = eval_preds
    predictions = np.argmax(logits, axis=-1)
    return metric.compute(predictions=predictions, references=labels)

trainer = Trainer(
    model,
    training_args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["validation"],
    data_collator=data_collator,
    tokenizer=tokenizer,
    compute_metrics=compute_metrics,
)


trainer.train()
#gr.Interface(
#    fn=trainer.train,
#    inputs=None,
#    outputs="training",
#    title="test",
#).launch()