Update README.md

README.md

@@ -37,46 +37,52 @@ license: cc-by-nc-4.0
 import torch
 from transformers import AutoModelForQuestionAnswering, AutoTokenizer
 
-context = "your context"
-question = "your question"
+def predict_answer(qa_text_pair):
+    # Encoding
+    encodings = tokenizer(
+        qa_text_pair['question'], qa_text_pair['context'], 
+        max_length=512, 
+        truncation=True,
+        padding="max_length", 
+        return_token_type_ids=False,
+        return_offsets_mapping=True
+    )
+    encodings = {key: torch.tensor([val]).to(device) for key, val in encodings.items()}
+    # Predict
+    with torch.no_grad():
+        pred = model(encodings['input_ids'], encodings['attention_mask'])
+        start_logits, end_logits = pred.start_logits, pred.end_logits
+        token_start_index, token_end_index = start_logits.argmax(dim=-1), end_logits.argmax(dim=-1)
+        pred_ids = encodings['input_ids'][0][token_start_index: token_end_index + 1]
+        # Answer start/end offset of context.
+        answer_start_offset = int(encodings['offset_mapping'][0][token_start_index][0][0])
+        answer_end_offset = int(encodings['offset_mapping'][0][token_end_index][0][1])
+        answer_offset = (answer_start_offset, answer_end_offset)
+    # Decoding
+    answer_text = tokenizer.decode(pred_ids) # text
+    del encodings 
+    return {'answer_text':answer_text, 'answer_offset':answer_offset}
 
-# Load fine-tuned MRC model by HuggingFace Model Hub
-HUGGINGFACE_MODEL_PATH = "bespin-global/klue-bert-base-aihub-mrc"
-tokenizer = AutoTokenizer.from_pretrained(HUGGINGFACE_MODEL_PATH )
-model = AutoModelForQuestionAnswering.from_pretrained(HUGGINGFACE_MODEL_PATH )
+
+# Load fine-tuned MRC model
+MODEL_PATH = "bespin-global/klue-bert-base-aihub-mrc"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
+model = AutoModelForQuestionAnswering.from_pretrained(MODEL_PATH)
 
 # gpu or cpu
 device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
 model.to(device)
 model.eval()
 
-# Encoding
-encodings = tokenizer(
-                context, question, 
-                max_length=512, 
-                truncation=True,
-                padding="max_length", 
-                return_token_type_ids=False,
-                return_offsets_mapping=True
-            )
-encodings = {key: torch.tensor([val]) for key, val in encodings.items()}             
-input_ids = encodings["input_ids"].to(device)
-attention_mask = encodings["attention_mask"].to(device)
-offset_mappings = encodings["offset_mapping"].to(device)
-
-# Predict
-pred = model(input_ids, attention_mask=attention_mask)
-start_logits, end_logits = pred.start_logits, pred.end_logits
-token_start_index, token_end_index = start_logits.argmax(dim=-1), end_logits.argmax(dim=-1)
-pred_ids = input_ids[0][token_start_index: token_end_index + 1]
-
-# Answer start/end offset of context.
-answer_start_offset = int(offset_mappings[0][token_start_index][0][0])
-answer_end_offset = int(offset_mappings[0][token_end_index][0][1])
+context = '''애플 M1(영어: Apple M1)은 애플이 자사의 매킨토시 컴퓨터용으로 설계한 최초의 ARM 기반 SoC이다. 
+4세대 맥북 에어, 5세대 맥 미니, 13인치 5세대 맥북 프로, 5세대 아이패드 프로에 선보였다. 5나노미터 공정을 사용하여 제조된 최초의 개인용 컴퓨터 칩이다. 
+애플은 저전력 실리콘의, 세계에서 가장 빠른 ARM 기반의 중앙 처리 장치(CPU) 코어, 그리고 세계 최고의 CPU 성능 대 와트를 갖추고 있다고 주장하고 있다.'''
+question = "애플이 m1에 대해 주장하는건 뭐야?"
 
-# Answer text
-answer_text = tokenizer.decode(pred_ids)
-print(f"ANSWER : {answer_text}")
+qa_text_pair = {'context':context, 'question':question}
+result = predict_answer(qa_text_pair)
+print('Answer Text: ', result['answer_text'])  # 저전력 실리콘의, 세계에서 가장 빠른 ARM 기반의 중앙 처리 장치 ( CPU ) 코어, 그리고 세계 최고의 CPU 성능 대 와트를 갖추고 있다고 주장하고 있다.
+print('Answer Offset: ', result['answer_offset'])  # (159, 246)
 ```
 
 ## Citing & Authors