Merge pull request #26 from dahongj/milestone-3

finetune.py

@@ -10,11 +10,13 @@ from transformers import DistilBertForSequenceClassification, AdamW
 
 model_name = "distilbert-base-uncased"
 
+#Reading text
 df = pd.read_csv('train.csv')
 train_texts = df["comment_text"].values
 train_labels = df[df.columns[2:]].values
 train_texts, val_texts, train_labels, val_labels = train_test_split(train_texts, train_labels, test_size=.2)
 
+#Dataset class to create the labels and encode them
 class TextDataset(Dataset):
   def __init__(self,texts,labels):
     self.texts = texts
@@ -30,21 +32,26 @@ class TextDataset(Dataset):
   def __len__(self):
     return len(self.labels)
 
+#This is the tokenizer for the current model
 tokenizer = DistilBertTokenizerFast.from_pretrained(model_name)
 
+#Set up the dataset
 train_dataset = TextDataset(train_texts,train_labels)
 val_dataset = TextDataset(val_texts, val_labels)
 
 device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
 
+#Use multilabel model because there are 6 variables to fintune for
 model = DistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased', num_labels=6, problem_type="multi_label_classification")
 model.to(device)
 model.train()
 
+#Use these parameters
 train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)
 
 optim = AdamW(model.parameters(), lr=5e-5)
 
+#Finetune process
 for epoch in range(1):
     for batch in train_loader:
         optim.zero_grad()
@@ -59,6 +66,8 @@ for epoch in range(1):
 
 model.eval()
 
+#Upload trained model to a file
 model.save_pretrained("sentiment_custom_model")
 
+#Upload tokenizer to a file
 tokenizer.save_pretrained("sentiment_tokenizer")