update submit

src/submission/submit.py

@@ -22,42 +22,50 @@ REQUESTED_MODELS = None
 USERS_TO_SUBMISSION_DATES = None
 
 def get_top_prediction(text, tokenizer, model):
-    inputs = tokenizer(text, return_tensors='pt')
-    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)
-        # 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']
-    option_logits = []
-    for option in options:
-        # 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:
-        print("No valid options found.")
-        return None
+    try:
+        inputs = tokenizer(text, return_tensors='pt')
+        if torch.cuda.is_available():
+            model = model.cuda()
+            inputs = {k: v.cuda() for k, v in inputs.items()}
+        else:
+            model = model.cpu()
 
-    # Get the option with the highest logit
-    top_option = max(option_logits, key=lambda x: x[0])[1]
-    return top_option
+        with torch.no_grad():
+            outputs = model(**inputs)
+            print(f"outputs.logits shape: {outputs.logits.shape}")
+            seq_len = outputs.logits.size(1)
+            if seq_len == 0:
+                print("No logits were produced by the model.")
+                return None
+            logits = outputs.logits[0, -1, :]  # Shape: [vocab_size]
+
+        options = ['A', 'B', 'C', 'D']
+        option_logits = []
+        for option in options:
+            # 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:
+            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
+    except Exception as e:
+        tb = traceback.format_exc()
+        print(f"Error in get_top_prediction: {e}\n{tb}")
+        return None
 
 
 def evaluate_model_accuracy(model_name, num_examples):