update submit

src/submission/submit.py

@@ -29,14 +29,24 @@ def get_top_prediction(text, tokenizer, model):
 
     with torch.no_grad():
         outputs = model(**inputs)
-        logits = outputs.logits[0, -1]
+        logits = outputs.logits[0, -1]  # Get logits of the last token
 
     options = [' A', ' B', ' C', ' D']
     option_logits = []
+
+    # Iterate through each option
     for option in options:
-        option_id = tokenizer(option).input_ids[-1]
-        option_logit = logits[option_id]
-        option_logits.append((option_logit.item(), option.strip()))
+        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.")
+
+    if not option_logits:
+        return "No valid options"
 
     # Get the option with the highest logit
     top_option = max(option_logits, key=lambda x: x[0])[1]