Upload code/inference.py with huggingface_hub

code/inference.py

@@ -0,0 +1,34 @@
+from sagemaker_inference import encoder
+import torch
+from transformers import AutoTokenizer, AutoModelForMultipleChoice
+
+
+def model_fn(model_dir):
+    tokenizer = AutoTokenizer.from_pretrained(model_dir)
+    model = AutoModelForMultipleChoice.from_pretrained(model_dir)
+    return {"model": model, "tokenizer": tokenizer}
+
+
+def predict_fn(data, model):
+    prompt = data["prompt"]
+    candidates = data["candidates"]
+
+    inputs = model["tokenizer"](
+        [[prompt, candidate] for candidate in candidates],
+        return_tensors="pt",
+        padding=True
+    )
+
+    labels = torch.tensor(0).unsqueeze(0)
+
+    with torch.no_grad():
+        outputs = model(
+            **{k: v.unsqueeze(0) for k, v in inputs.items()}, labels=labels
+        )
+    
+    return outputs.logits
+
+
+def output_fn(prediction, content_type):
+        result = {i: x for i, x in enumerate(prediction)}
+        return encoder.encode(result, content_type)