How to use
import torch
from transformers import AutoTokenizer, AutoModelForQuestionAnswering
MODEL_ID = "jidoren/mmbert-squad-v2"
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForQuestionAnswering.from_pretrained(MODEL_ID)
model.eval()
question = "Where was Einstein born?"
context = """
Albert Einstein was born in Ulm, Germany, in 1879.
He developed the theory of relativity.
"""
inputs = tokenizer(
question,
context,
return_tensors="pt",
truncation=True,
)
with torch.no_grad():
outputs = model(**inputs)
start = outputs.start_logits.argmax()
end = outputs.end_logits.argmax()
answer = tokenizer.decode(
inputs["input_ids"][0][start:end + 1],
skip_special_tokens=True,
)
print(answer)
