update roberta script

2021-06-22 14:22:05 +02:00 · 2021-06-22 14:22:05 +02:00 · e348d16dde
parent fb4b0d95e3
commit e348d16dde
3 changed files with 1413 additions and 1411 deletions
--- a/dev-0/out.tsv
+++ b/dev-0/out.tsv
--- a/roberta.py
+++ b/roberta.py
@ -34,6 +34,8 @@ class IMDbDataset(torch.utils.data.Dataset):
 def prepare(data_train_X, data_train_Y):
    tokenizer = AutoTokenizer.from_pretrained(PRE_TRAINED)
    model = AutoModelForSequenceClassification.from_pretrained(PRE_TRAINED, num_labels=2)
+    device = torch.device("cpu")
+    model.to(device)
    encoded_input  = tokenizer([text[0] for text in list(zip(data_train_X, data_train_Y))], truncation=True, padding=True)
    train_dataset = IMDbDataset(encoded_input , [int(text[1]) for text in list(zip(data_train_X, data_train_Y))])

--- a/test-A/out.tsv
+++ b/test-A/out.tsv