| 31 | return data |
| 32 | # post process |
| 33 | def get_capital_answer(text): |
| 34 | patterns = [ |
| 35 | "the answer is ([A-E])", |
| 36 | "the answer is([A-E])", |
| 37 | "Answer: ([A-E])", |
| 38 | "Answer: \(([A-E])\)", |
| 39 | "Option \(([A-E])\)", |
| 40 | "Answer:([A-E])", |
| 41 | "Option ([A-E])", |
| 42 | "Opt ([A-E])" |
| 43 | ] |
| 44 | for pattern in patterns: |
| 45 | match = re.search(pattern,text,re.IGNORECASE) |
| 46 | if match: |
| 47 | return match.group(1) |
| 48 | match = re.findall("[A-D]", text) |
| 49 | if match: |
| 50 | return match[0] |
| 51 | return "" |
| 52 | if __name__ == "__main__": |
| 53 | tokenizer = AutoTokenizer.from_pretrained(PATH) |
| 54 | model = AutoModelForCausalLM.from_pretrained(PATH, trust_remote_code=True, device_map="auto", |