| 32 | |
| 33 | |
| 34 | class GeminiInference: |
| 35 | def __init__(self): |
| 36 | self.model = genai.GenerativeModel( |
| 37 | model_name="gemini-1.5-pro-latest", |
| 38 | safety_settings=safety_settings, |
| 39 | generation_config=generation_config, |
| 40 | ) |
| 41 | |
| 42 | def post_process(self, text): |
| 43 | match = re.search(r"(?i)(?<=\banswer:\s).*", text) |
| 44 | if match: |
| 45 | return match.group(0) |
| 46 | else: |
| 47 | return text |
| 48 | |
| 49 | def predict(self, prompt): |
| 50 | chat_session = self.model.start_chat( |
| 51 | history=[] |
| 52 | ) |
| 53 | response = chat_session.send_message(prompt) |
| 54 | return self.post_process(response.text) |
| 55 | |
| 56 | def predict_nq(self, context, question, titles): |
| 57 | prompt = (f"Go through the following context and then extract the answer of the question from the context. " |
| 58 | f"The context is a list of Wikipedia documents, ordered by title: {titles}. " |
| 59 | f"Each Wikipedia document contains a title field and a text field. " |
| 60 | f"The context is: {context}. " |
| 61 | f"Find the useful documents from the context, then extract the answer to answer the question: {question}." |
| 62 | f"Answer the question directly. Your response should be very concise. ") |
| 63 | long_answer = self.predict(prompt) |
| 64 | short_answer = self.extract_answer(question, long_answer) |
| 65 | return long_answer, short_answer |
| 66 | |
| 67 | def predict_hotpotqa(self, context, question, titles): |
| 68 | prompt = (f"Go through the following context and then answer the question " |
| 69 | f"The context is a list of Wikipedia documents titled: {titles}. " |
| 70 | f"There are two types of questions: comparison questions, which require a yes or no answer or a selection from two candidates, " |
| 71 | f"and general questions, which demand a concise response. " |
| 72 | f"The context is: {context}. " |
| 73 | f"Find the useful documents from the context, then answer the question: {question}." |
| 74 | f"For general questions, you should use the exact words from the context as the answer to avoid ambiguity. " |
| 75 | f"Answer the question directly and don't output other thing. ") |
| 76 | long_answer = self.predict(prompt) |
| 77 | short_answer = self.extract_answer(question, long_answer) |
| 78 | return long_answer, short_answer |
| 79 | |
| 80 | def predict_close_book(self, question, demo_file_path, num_demo=16): |
| 81 | demo = load_json_file(demo_file_path) |
| 82 | prompt = ("Here are some examples of questions and their corresponding answer, each with a 'Question' field and an 'Answer' field. " |
| 83 | "Answer the question directly and don't output other thing. ") |
| 84 | for item in demo[:num_demo]: |
| 85 | prompt += f"Question: {item['question']} Answer: {item['short_answers'][0]}\n" |
| 86 | prompt += f"Question: {question} Answer: " |
| 87 | answer = self.predict(prompt) |
| 88 | return answer |
| 89 | |
| 90 | def generate_demo_examples(self, num_demo=4): |
| 91 | if num_demo == 0: |