Added comments and made a start on infer.py

src/infer.py

@@ -2,14 +2,21 @@ import pandas as pd
 import numpy as np
 import torch
 import argparse
-from transformers import AutoTokenizer
 from torch.utils.tensorboard import SummaryWriter
+from transformers import AutoTokenizer
+from transformers import AutoModelForSequenceClassification
 
 # mappings
 binary_map = {1:'Yes', 0:'No'}
 aspect_map = {0:'App', 1:'Driver', 2:'General', 3:'Payment', 4:'Pricing', 5:'Service'}
 sentiment_map = {0:'Positive', 1:'Neutral', 2:'Negative'}
 
+label_names = {
+    'bug_report': ['No', 'Yes'],
+    'feature_request': ['No', 'Yes'],
+    'aspect': ['App', 'Driver', 'General', 'Payment', 'Pricing', 'Service'],
+    'aspect_sentiment': ['Positive', 'Neutral', 'Negative']
+}
 
 SEED = 4321
 torch.manual_seed(SEED)
@@ -17,9 +24,31 @@ np.random.seed(SEED)
 
 def parse_args():
     parser = argparse.ArgumentParser(description="RECLASS, Multitask learning for review classification.")
-    parser.add_argument("--model_path", type=str, help="Enter the models path / the desired .pt file")
-    parser.add_argument("--task", type=str, default="all", choices=["all", "bug_report", "feature_request", "aspect", "aspect_sentiment"], help="Specific task to train for stl usage only" )
-    parser.add_argument("--interactive", help="Loops reading input until exit")
+    parser.add_argument("--model_path", type=str, required=True, help=".pt file path")
+    parser.add_argument("--task", type=str, default="all", choices=["all", "bug_report", "feature_request", "aspect", "aspect_sentiment"])
+    parser.add_argument("--interactive", help="Loops reading input until exit()")
     parser.add_argument("--text", help="Use command line text for input")
+    parser.add_argument("--dataset", type=str, required=True, help="Enter a file for inference")
+
     return parser.parse_args()
 
+
+
+def main():
+    args = parse_args()
+    print(f'='*50)
+    print(f' '*15 + "Starting inference")
+    if torch.cuda.is_available():
+        print(f' '*15 + "GPU:", torch.cuda.get_device_name(0))
+        torch.cuda.manual_seed_all(SEED)
+        torch.cuda.manual_seed(SEED)
+    else:
+        print(f' '*15 + "No GPUs available")
+    print(f'='*50 + "\n")
+    print(f"Running inference on: {args.model_path.upper()} using {args.dataset}")
+
+    tokenizer = AutoTokenizer.from_pretrained("FacebookAI/xlm-roberta-base")
+    infer_data = f"data/processed/{args.dataset}_infer.csv"
+
+if __name__ == main():
+    main()

src/model.py

@@ -52,7 +52,7 @@ class Model(nn.Module):
 
         # Applied across shared cls token, before all task heads 
         self.dropout = nn.Dropout(dropout_rate)
-
+        # get logits for each head
         self.bug_head = nn.Linear(hidden_size, 2)
         self.feature_head = nn.Linear(hidden_size, 2)
         self.aspect_head = nn.Linear(hidden_size, 6)

src/preprocess.py

@@ -160,8 +160,8 @@ def preprocess_uber_reviews(input_path, output_path):
     return df_clean
 
 if __name__ == "__main__":
-    input_file = "multitag/data/uber_reviews.csv"
-    output_file = "multitag/data/uber_reviews_cleaned.csv"
+    input_file = "data/raw/uber_reviews.csv"
+    output_file = "data/raw/uber_reviews_cleaned.csv"
     
     df_clean = preprocess_uber_reviews(input_file, output_file)
     print("\nPreprocessing complete!")

src/sampler.py

@@ -190,7 +190,6 @@ class Sampler:
         mini_sample = self.data.sample(200)     #   reading some samples manually
         return mini_sample
 
-         
     
     def save_sample(self, sample_df,output_path):
         """Save sample and display statistics"""

src/train.py

@@ -126,6 +126,7 @@ def main():
     print("Aspect sentiment class weights:", aspect_sentiment_weights.cpu().numpy())
     
     # equal weighted task losses. unequal was considered but equal weights performed well without adding complexity
+    # CrossEntropyLoss = LogSoftmax + NLLLoss (negative log likelihood) 
     criterions = {
         'bug_report': nn.CrossEntropyLoss(weight=bug_weights),
         'feature_request': nn.CrossEntropyLoss(weight=feature_weights),
@@ -134,6 +135,7 @@ def main():
     }
 
     # -------------------- Optimizer and scheduler -------------------
+    # adaptive momentum and weight decay keeps track of previous weight adaptions and ensures they dont get too large (weight also shrinks towards 0 each pass) 
     optimizer = torch.optim.AdamW(
         model.parameters(), 
         lr=args.lr,