import torch from transformers import AutoTokenizer, AutoModelForQuestionAnswering
import tokenizer
tokenizer = AutoTokenizer.from_pretrained("atharvamundada99/bert-large-question-answering-finetuned-legal",cache_dir="/E/HUG_Models") model = AutoModelForQuestionAnswering.from_pretrained("atharvamundada99/bert-large-question-answering-finetuned-legal", cache_dir="/E/HUG_Models")
def get_answer( question, context): inputs = tokenizer(question, context, return_tensors="pt") with torch.no_grad(): outputs = model(**inputs) answer_start_index = outputs.start_logits.argmax() answer_end_index = outputs.end_logits.argmax()
predict_answer_tokens = inputs.input_ids[0, answer_start_index: answer_end_index + 1]
answer=tokenizer.decode(predict_answer_tokens, skip_special_tokens=True)
return answer
print(get_answer("What is your name", "My name is JACK")) #Output JACK