grady / util.py
bstraehle's picture
Update util.py
4dc6d7a verified
raw
history blame
2.68 kB
import os
import pandas as pd
from google import genai
def get_questions(file_path):
df = pd.read_json(file_path, lines=True)
result=[]
for index, row in df.iterrows():
result.append([row["Level"], row["Question"], row["file_name"], row["Final answer"]])
return result
def read_file(file_path):
df = None
ext = os.path.splitext(file_path)[1].lower()
if ext == ".csv":
df = pd.read_csv(file_path)
elif ext in (".xls", ".xlsx"):
df = pd.read_excel(file_path)
elif ext in (".json", ".jsonl"):
df = pd.read_json(file_path)
return "" if df is None else df.to_json()
def get_final_answer(model, question, answer):
prompt_template = """
You are an expert question answering assistant. Given a question and an initial answer, your task is to provide the final answer.
Your final answer must be a number and/or string OR as few words as possible OR a comma-separated list of numbers and/or strings.
If you are asked for a number, don't use comma to write your number neither use units such as USD, $, percent, or % unless specified otherwise.
If you are asked for a string, don't use articles, neither abbreviations (for example cities), and write the digits in plain text unless specified otherwise.
If you are asked for a comma-separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
If the final answer is a number, use a number not a word.
If the final answer is a word, start with an uppercase character.
If the final answer is a comma-separated list of numbers, use a space character after each comma.
If the final answer is a comma-separated list of strings, use a space character after each comma and start with a lowercase character.
Do not add any content to the final answer that is not in the initial answer.
**Question:** """ + question + """
**Initial answer:** """ + answer + """
**Example 1:** How many 'r's are in strawberry? 3
**Example 2:** What is the opposite of black? White
**Example 3:** What is the biggest city in California? Los Angeles
**Example 4:** What is the superlative of good? Best
**Example 5:** What are the first 10 numbers in the Fibonacci sequence? 0, 1, 1, 2, 3, 5, 8, 13, 21, 34
**Final answer:**
"""
client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
response = client.models.generate_content(
model=model,
contents=[prompt_template]
)
return response.text