Update README.md

README.md

@@ -71,6 +71,30 @@ model_path = "real-jiakai/roberta-base-uncased-finetuned-swag"
 tokenizer = AutoTokenizer.from_pretrained(model_path)
 model = AutoModelForMultipleChoice.from_pretrained(model_path)
 
+def predict_swag(context, endings, model, tokenizer):
+    encoding = tokenizer(
+        [context] * 4,
+        endings,
+        truncation=True,
+        max_length=128,
+        padding="max_length",
+        return_tensors="pt"
+    )
+    
+    input_ids = encoding['input_ids'].unsqueeze(0)
+    attention_mask = encoding['attention_mask'].unsqueeze(0)
+    
+    outputs = model(input_ids=input_ids, attention_mask=attention_mask)
+    logits = outputs.logits
+    
+    predicted_idx = torch.argmax(logits).item()
+    
+    return {
+        'context': context,
+        'predicted_ending': endings[predicted_idx],
+        'probabilities': torch.softmax(logits, dim=1)[0].tolist()
+    }
+
 # Example scenarios
 test_examples = [
     {
@@ -93,29 +117,22 @@ test_examples = [
     }
 ]
 
-def predict_swag(context, endings, model, tokenizer):
-    encoding = tokenizer(
-        [context] * 4,
-        endings,
-        truncation=True,
-        max_length=128,
-        padding="max_length",
-        return_tensors="pt"
+for i, example in enumerate(test_examples, 1):
+    result = predict_swag(
+        example['context'],
+        example['endings'],
+        model,
+        tokenizer
     )
     
-    input_ids = encoding['input_ids'].unsqueeze(0)
-    attention_mask = encoding['attention_mask'].unsqueeze(0)
-    
-    outputs = model(input_ids=input_ids, attention_mask=attention_mask)
-    logits = outputs.logits
-    
-    predicted_idx = torch.argmax(logits).item()
-    
-    return {
-        'context': context,
-        'predicted_ending': endings[predicted_idx],
-        'probabilities': torch.softmax(logits, dim=1)[0].tolist()
-    }
+    print(f"\n=== Test Scenario {i} ===")
+    print(f"Initial Context: {result['context']}")
+    print(f"\nPredicted Most Likely Ending: {result['predicted_ending']}")
+    print("\nProbabilities for All Options:")
+    for idx, (ending, prob) in enumerate(zip(result['all_endings'], result['probabilities'])):
+        print(f"Option {idx}: {ending}")
+        print(f"Probability: {prob:.3f}")
+    print("\n" + "="*50)
 ```
 
 ## Limitations and Biases