update submit

src/submission/submit.py

@@ -26,32 +26,40 @@ def get_top_prediction(text, tokenizer, model):
     if torch.cuda.is_available():
         model = model.cuda()
         inputs = {k: v.cuda() for k, v in inputs.items()}
+    else:
+        model = model.cpu()
 
     with torch.no_grad():
         outputs = model(**inputs)
-        logits = outputs.logits[0, -1]  # Get logits of the last token
+        # outputs.logits shape: [batch_size, seq_len, vocab_size]
+        # We want the logits for the last token
+        logits = outputs.logits[0, -1, :]  # Shape: [vocab_size]
 
-    options = [' A', ' B', ' C', ' D']
+    options = ['A', 'B', 'C', 'D']
     option_logits = []
-
-    # Iterate through each option
     for option in options:
-        option_ids = tokenizer(option).input_ids
-        # Ensure option_ids are within range and not empty
-        if option_ids and option_ids[-1] < logits.size(0):
-            option_id = option_ids[-1]
-            option_logit = logits[option_id]
-            option_logits.append((option_logit.item(), option.strip()))
-        else:
-            print(f"Skipping option '{option}' due to index out of range.")
+        # Encode the option without adding special tokens
+        option_ids = tokenizer.encode(option, add_special_tokens=False)
+        if not option_ids:
+            print(f"Option '{option}' could not be tokenized.")
+            continue
+        option_id = option_ids[0]
+        vocab_size = logits.size(0)
+        if option_id >= vocab_size:
+            print(f"Option ID {option_id} is out of bounds for vocabulary size {vocab_size}")
+            continue
+        option_logit = logits[option_id]
+        option_logits.append((option_logit.item(), option))
 
     if not option_logits:
-        return "No valid options"
+        print("No valid options found.")
+        return None
 
     # Get the option with the highest logit
     top_option = max(option_logits, key=lambda x: x[0])[1]
     return top_option
 
+
 def evaluate_model_accuracy(model_name, num_examples):
     try:
         # Load the model and tokenizer