| 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 |