train and validation data stored. train.py now can generate models

data/train.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6b9de2fb205c4a8fe3d082ae0441872f983ed8b07ae4bb965c7cab2822ecc453
-size 58540950
+oid sha256:bbe06c7232f904c0501c4dc5b950e4243d887cd7ffab04bfeaa12732514b47c8
+size 58602006

data/val.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:cea25d0297302c68eccc08d662d03b35af3f99106a2e0854a44ceace0a9204a1
-size 32710564
+oid sha256:d25635377a60d308197c5a8c0c7df7575953d29cfab91b6c11316266c6a5b27c
+size 32744803

src/train.py

@@ -4,11 +4,16 @@ import numpy as np
 import os
 import json
 import torch
+import sys
+
 from torch.utils.data import Dataset, DataLoader
 from transformers import DistilBertTokenizerFast, DistilBertForSequenceClassification
 from transformers import Trainer, TrainingArguments, AdamW
 
+torch.backends.cuda.matmul.allow_tf32 = True
 model_name = "distilbert-base-uncased"
+upsto_abstracts_model_path = './models/upsto_abstracts'
+upsto_claims_model_path = './models/upsto_claims'
 
 class USPTODataset(Dataset):
     def __init__(self, encodings, labels):
@@ -67,30 +72,26 @@ def LoadDataset():
     trainDF2 = trainDF.replace({"decision": yKey})
     valDF2 = valDF.replace({"decision": yKey})
 
-    # We combine the `abstract` and `claims` columns into a single `text` column. 
-    # We also re-label the `decision` column to `label`.
-    print("Combining columns and renaming `decision` to `label`")
+    # We re-label the `decision` column to `label`.
+    print("Renaming `decision` to `label`")
     trainDF3 = trainDF2.rename(columns={'decision': 'label'})
-    trainDF3['text'] = trainDF3['abstract'] + ' ' + trainDF3['claims']
-    trainDF3.drop(columns=["abstract","claims"],inplace=True)
-
     valDF3 = valDF2.rename(columns={'decision': 'label'})
-    valDF3['text'] = valDF3['abstract'] + ' ' + valDF3['claims']
-    valDF3.drop(columns=["abstract","claims"],inplace=True)
 
     # We can grab the data for each column so that we have a list of values for training labels, 
     #   training texts, validation labels, and validation texts.
     print("Extracting label and text data from dataframes")
     trainData = {
         "labels":trainDF3["label"].tolist(),
-        "text":trainDF3["text"].tolist()
+        "abstracts":trainDF3["abstract"].tolist(),
+        "claims":trainDF3["claims"].tolist(),
     }
     valData = {
         "labels":valDF3["label"].tolist(),
-        "text":valDF3["text"].tolist()
+        "abstracts":valDF3["abstract"].tolist(),
+        "claims":valDF3["claims"].tolist(),
     }
-    print(f'TRAINING:\t# labels: {len(trainData["labels"])}\t# texts: {len(trainData["text"])}')
-    print(f'VALID:\t# labels: {len(valData["labels"])}\t# texts: {len(valData["text"])}')
+    #print(f'TRAINING:\t# labels: {len(trainData["labels"])}\t# texts: {len(trainData["text"])}')
+    #print(f'VALID:\t# labels: {len(valData["labels"])}\t# texts: {len(valData["text"])}')
 
     if not os.path.exists("./data"):
         os.makedirs('./data')
@@ -102,6 +103,79 @@ def LoadDataset():
 
     return trainData, valData
 
+def TrainModel(trainData, valData):
+    print("=== ENCODING DATA ===")
+    #print(len(trainData["labels"]), len(trainData["text"]), len(valData["labels"]), len(valData["text"]))
+    print("\t- initializing tokenizer")
+    tokenizer = DistilBertTokenizerFast.from_pretrained(model_name)
+    print("\t- encoding training data")
+    train_abstracts_encodings = tokenizer(trainData["abstracts"], truncation=True, padding=True)
+    train_claims_encodings = tokenizer(trainData["claims"], truncation=True, padding=True)
+    #print("\t- encoding validation data")
+    #val_abstracts_encodings = tokenizer(valData["abstracts"], truncation=True, padding=True)
+    #val_claims_encodings = tokenizer(valData["claims"], truncation=True, padding=True)
+
+    print("=== CREATING DATASETS ===")
+    print("\t- initializing dataset for training data")
+    train_abstracts_dataset = USPTODataset(train_abstracts_encodings, trainData["labels"])
+    train_claims_dataset = USPTODataset(train_claims_encodings, trainData["labels"])
+    #print("\t- initializing dataset for validation data")
+    #val_abstracts_dataset = USPTODataset(val_abstracts_encodings, valData["labels"])
+    #val_claims_dataset = USPTODataset(val_claims_encodings, valData["labels"])
+
+    print("=== PREPARING MODEL ===")
+    print("\t- setting up device")
+    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+    print("\t- initializing model")
+    model = DistilBertForSequenceClassification.from_pretrained(model_name)
+    model.to(device)
+    model.train()
+
+    print("== PREPARING TRAINING ===")
+    print("\t- initializing trainers")
+    train_abstracts_loader = DataLoader(train_abstracts_dataset, batch_size=4, shuffle=True)
+    train_claims_loader = DataLoader(train_claims_dataset, batch_size=4, shuffle=True)
+    #train_claims_loader = DataLoader(train_claims_dataset, batch_size=4, shuffle=True)
+    print("\t- initializing optim")
+    optim = AdamW(model.parameters(), lr=5e-5)
+
+    def Train(loader, save_path, num_train_epochs=2):
+        batch_num = len(loader)
+        for epoch in range(num_train_epochs):
+            print(f'\t- Training epoch {epoch+1}/{num_train_epochs}')
+            batch_count = 0
+            for batch in loader:
+                print(f'{batch_count}|{batch_num} - {round((batch_count/batch_num)*100)}%', end="")
+                #print('\t\t- optim zero grad')
+                optim.zero_grad()
+                #print('\t\t- input_ids')
+                input_ids = batch['input_ids'].to(device)
+                #print('\t\t- attention_mask')
+                attention_mask = batch['attention_mask'].to(device)
+                #print('\t\t- labels0')
+                labels = batch['labels'].to(device)
+                #print('\t\t- outputs')
+                outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
+            
+                #print('\t\t- loss')
+                loss = outputs[0]
+                #print('\t\t- backwards')
+                loss.backward()
+                #print('\t\t- step')
+                optim.step()
+
+                batch_count += 1
+                print("\r", end="")
+    
+        model.eval()
+        model.save_pretrained(save_path, from_pt=True) 
+        print(f'Saved model in {save_path}!')
+    
+    print("=== TRAINING ABSTRACTS ===")
+    Train(train_abstracts_loader,upsto_abstracts_model_path)
+    print("=== TRAINING CLAIMS ===")
+    Train(train_claims_loader,upsto_claims_model_path)
+
 def main():
     trainDataPath = "./data/train.json"
     valDataPath = "./data/val.json"
@@ -109,6 +183,7 @@ def main():
     valData = None
 
     if os.path.exists(trainDataPath) and os.path.exists(valDataPath):
+        print("Loading from existing data files")
         ftrain = open(trainDataPath)
         trainData = json.load(ftrain)
         ftrain.close()
@@ -118,15 +193,12 @@ def main():
     else:
         trainData, valData = LoadDataset()
     
-    print(len(trainData["labels"]), len(trainData["text"]), len(valData["labels"]), len(valData["text"]))
-
-    tokenizer = DistilBertTokenizerFast.from_pretrained(model_name)
-    train_encodings = tokenizer(trainData["text"], truncation=True, padding=True)
-    val_encodings = tokenizer(valData["text"], truncation=True, padding=True)
+    #print(len(trainData["labels"]), len(trainData["text"]), len(valData["labels"]), len(valData["text"]))
+    print("Data loaded successfully!")
 
-    train_dataset = USPTODataset(train_encodings, trainData["labels"])
-    val_dataset = USPTODataset(val_encodings, valData["labels"])
+    TrainModel(trainData, valData)
 
+    """
     train_args = TrainingArguments(
         output_dir="./results",
         num_train_epochs=2,
@@ -147,6 +219,7 @@ def main():
         eval_dataset=val_dataset
     )
     trainer.train()
+    """
 
 if __name__ == "__main__":
     main()