AutoRAG 0.0.0__py3-none-any.whl

autorag/data/qa/generation_gt/openai_gen_gt.py

@@ -0,0 +1,84 @@
+import itertools
+from typing import Dict
+
+from openai import AsyncClient
+from pydantic import BaseModel
+
+from autorag.data.qa.generation_gt.base import add_gen_gt
+from autorag.data.qa.generation_gt.prompt import GEN_GT_SYSTEM_PROMPT
+
+
+class Response(BaseModel):
+	answer: str
+
+
+async def make_gen_gt_openai(
+	row: Dict,
+	client: AsyncClient,
+	system_prompt: str,
+	model_name: str = "gpt-4o-2024-08-06",
+):
+	retrieval_gt_contents = list(
+		itertools.chain.from_iterable(row["retrieval_gt_contents"])
+	)
+	query = row["query"]
+	passage_str = "\n".join(retrieval_gt_contents)
+	user_prompt = f"Text:\n<|text_start|>\n{passage_str}\n<|text_end|>\n\nQuestion:\n{query}\n\nAnswer:"
+
+	completion = await client.beta.chat.completions.parse(
+		model=model_name,
+		messages=[
+			{"role": "system", "content": system_prompt},
+			{"role": "user", "content": user_prompt},
+		],
+		temperature=0.0,
+		response_format=Response,
+	)
+	response: Response = completion.choices[0].message.parsed
+	return add_gen_gt(row, response.answer)
+
+
+async def make_concise_gen_gt(
+	row: Dict,
+	client: AsyncClient,
+	model_name: str = "gpt-4o-2024-08-06",
+	lang: str = "en",
+):
+	"""
+	Generate concise generation_gt using OpenAI Structured Output for preventing errors.
+	It generates a concise answer, so it is generally a word or just a phrase.
+
+	:param row: The input row of the qa dataframe.
+	:param client: The OpenAI async client.
+	:param model_name: The model name that supports structured output.
+	    It has to be "gpt-4o-2024-08-06" or "gpt-4o-mini-2024-07-18".
+	:param lang: The language code of the prompt.
+		Default is "en".
+	:return: The output row of the qa dataframe with added "generation_gt" in it.
+	"""
+	return await make_gen_gt_openai(
+		row, client, GEN_GT_SYSTEM_PROMPT["concise"][lang], model_name
+	)
+
+
+async def make_basic_gen_gt(
+	row: Dict,
+	client: AsyncClient,
+	model_name: str = "gpt-4o-2024-08-06",
+	lang: str = "en",
+):
+	"""
+	Generate basic generation_gt using OpenAI Structured Output for preventing errors.
+	It generates a "basic" answer, and its prompt is simple.
+
+	:param row: The input row of the qa dataframe.
+	:param client: The OpenAI async client.
+	:param model_name: The model name that supports structured output.
+	    It has to be "gpt-4o-2024-08-06" or "gpt-4o-mini-2024-07-18".
+	:param lang: The language code of the prompt.
+		Default is "en".
+	:return: The output row of the qa dataframe with added "generation_gt" in it.
+	"""
+	return await make_gen_gt_openai(
+		row, client, GEN_GT_SYSTEM_PROMPT["basic"][lang], model_name
+	)

autorag/data/qa/generation_gt/prompt.py

@@ -0,0 +1,27 @@
+GEN_GT_SYSTEM_PROMPT = {
+	"concise": {
+		"en": """You are an AI assistant to answer the given question in the provide evidence text.
+You can find the evidence from the given text about question, and you have to write a proper answer to the given question.
+Your answer have to be concise and relevant to the question.
+Do not make a verbose answer and make it super clear.
+It doesn't have to be an full sentence. It can be the answer is a word or a paraphrase.""",
+		"ko": """당신은 주어진 질문에 대해 제공된 Text 내에서 답을 찾는 AI 비서입니다.
+질문에 대한 답을 Text에서 찾아 적절한 답변을 작성하세요.
+답변은 간결하고 질문에 관련된 내용만 포함해야 합니다.
+불필요하게 길게 답변하지 말고, 명확하게 작성하세요.
+완전한 문장이 아니어도 되며, 답은 단어나 요약일 수 있습니다.""",
+		"ja": """
+あなたは与えられた質問に対して提供されたText内で答えを探すAI秘書です。
+質問に対する答えをTextで探して適切な答えを作成しましょう。
+回答は簡潔で、質問に関連する内容のみを含める必要があります。
+不必要に長く答えず、明確に作成しましょう。
+完全な文章でなくてもいいし、答えは単語や要約かもしれません。
+""",
+	},
+	"basic": {
+		"en": """You are an AI assistant to answer the given question in the provide evidence text.
+You can find the evidence from the given text about question, and you have to write a proper answer to the given question.""",
+		"ko": "당신은 주어진 질문에 대한 답을 제공된 Text 내에서 찾는 AI 비서입니다. 질문과 관련된 증거를 Text에서 찾아 적절한 답변을 작성하세요.",
+		"ja": "あなたは与えられた質問に対する答えを提供されたText内で探すAI秘書です。 質問に関する証拠をTextで探して適切な回答を作成しましょう。",
+	},
+}

autorag/data/qa/query/llama_gen_query.py

@@ -0,0 +1,82 @@
+import itertools
+from typing import Dict, List
+
+from llama_index.core.base.llms.base import BaseLLM
+from llama_index.core.base.llms.types import ChatResponse, ChatMessage, MessageRole
+
+from autorag.data.qa.query.prompt import QUERY_GEN_PROMPT, QUERY_GEN_PROMPT_EXTRA
+
+
+async def llama_index_generate_base(
+	row: Dict,
+	llm: BaseLLM,
+	messages: List[ChatMessage],
+) -> Dict:
+	context = list(itertools.chain.from_iterable(row["retrieval_gt_contents"]))
+	context_str = "\n".join([f"{i + 1}. {c}" for i, c in enumerate(context)])
+	user_prompt = f"Text:\n{context_str}\n\nGenerated Question from the Text:\n"
+	user_message = ChatMessage(role=MessageRole.USER, content=user_prompt)
+	new_messages = [*messages, user_message]
+	chat_response: ChatResponse = await llm.achat(messages=new_messages)
+	row["query"] = chat_response.message.content
+	return row
+
+
+async def factoid_query_gen(
+	row: Dict,
+	llm: BaseLLM,
+	lang: str = "en",
+) -> Dict:
+	return await llama_index_generate_base(
+		row, llm, QUERY_GEN_PROMPT["factoid_single_hop"][lang]
+	)
+
+
+async def concept_completion_query_gen(
+	row: Dict,
+	llm: BaseLLM,
+	lang: str = "en",
+) -> Dict:
+	return await llama_index_generate_base(
+		row, llm, QUERY_GEN_PROMPT["concept_completion"][lang]
+	)
+
+
+async def two_hop_incremental(
+	row: Dict,
+	llm: BaseLLM,
+	lang: str = "en",
+) -> Dict:
+	messages = QUERY_GEN_PROMPT["two_hop_incremental"][lang]
+	passages = row["retrieval_gt_contents"]
+	assert len(passages) >= 2, (
+		"You have to sample more than two passages for making two-hop questions."
+	)
+	context_str = f"Document 1: {passages[0][0]}\nDocument 2: {passages[1][0]}"
+	user_prompt = f"{context_str}\n\nGenerated two-hop Question from two Documents:\n"
+	messages.append(ChatMessage(role=MessageRole.USER, content=user_prompt))
+
+	chat_response: ChatResponse = await llm.achat(messages=messages)
+	response = chat_response.message.content
+	row["query"] = response.split(":")[-1].strip()
+	return row
+
+
+async def custom_query_gen(
+	row: Dict,
+	llm: BaseLLM,
+	messages: List[ChatMessage],
+) -> Dict:
+	return await llama_index_generate_base(row, llm, messages)
+
+
+# Experimental feature: can only use factoid_single_hop
+async def multiple_queries_gen(
+	row: Dict,
+	llm: BaseLLM,
+	lang: str = "en",
+	n: int = 3,
+) -> Dict:
+	_messages = QUERY_GEN_PROMPT["factoid_single_hop"][lang]
+	_messages[0].content += QUERY_GEN_PROMPT_EXTRA["multiple_queries"][lang].format(n=n)
+	return await llama_index_generate_base(row, llm, _messages)

autorag/data/qa/query/openai_gen_query.py

@@ -0,0 +1,95 @@
+import itertools
+from typing import Dict, List
+
+from llama_index.core.base.llms.types import ChatMessage, MessageRole
+from llama_index.llms.openai.utils import to_openai_message_dicts
+from openai import AsyncClient
+from pydantic import BaseModel
+
+from autorag.data.qa.query.prompt import QUERY_GEN_PROMPT
+
+
+class Response(BaseModel):
+	query: str
+
+
+# Single hop QA generation OpenAI
+async def query_gen_openai_base(
+	row: Dict,
+	client: AsyncClient,
+	messages: List[ChatMessage],
+	model_name: str = "gpt-4o-2024-08-06",
+):
+	context = list(itertools.chain.from_iterable(row["retrieval_gt_contents"]))
+	context_str = "Text:\n" + "\n".join(
+		[f"{i + 1}. {c}" for i, c in enumerate(context)]
+	)
+	user_prompt = f"{context_str}\n\nGenerated Question from the Text:\n"
+	messages.append(ChatMessage(role=MessageRole.USER, content=user_prompt))
+
+	completion = await client.beta.chat.completions.parse(
+		model=model_name,
+		messages=to_openai_message_dicts(messages),
+		response_format=Response,
+	)
+	row["query"] = completion.choices[0].message.parsed.query
+	return row
+
+
+async def factoid_query_gen(
+	row: Dict,
+	client: AsyncClient,
+	model_name: str = "gpt-4o-2024-08-06",
+	lang: str = "en",
+) -> Dict:
+	return await query_gen_openai_base(
+		row, client, QUERY_GEN_PROMPT["factoid_single_hop"][lang], model_name
+	)
+
+
+async def concept_completion_query_gen(
+	row: Dict,
+	client: AsyncClient,
+	model_name: str = "gpt-4o-2024-08-06",
+	lang: str = "en",
+) -> Dict:
+	return await query_gen_openai_base(
+		row, client, QUERY_GEN_PROMPT["factoid_single_hop"][lang], model_name
+	)
+
+
+class TwoHopIncrementalResponse(BaseModel):
+	answer: str
+	one_hop_question: str
+	two_hop_question: str
+
+
+async def two_hop_incremental(
+	row: Dict,
+	client: AsyncClient,
+	model_name: str = "gpt-4o-2024-08-06",
+	lang: str = "en",
+) -> Dict:
+	"""
+	Create a two-hop question using incremental prompt.
+	Incremental prompt is more effective to create multi-hop question.
+	The input retrieval_gt has to include more than one passage.
+
+	:return: The two-hop question using openai incremental prompt
+	"""
+	messages = QUERY_GEN_PROMPT["two_hop_incremental"][lang]
+	passages = row["retrieval_gt_contents"]
+	assert len(passages) >= 2, (
+		"You have to sample more than two passages for making two-hop questions."
+	)
+	context_str = f"Document 1: {passages[0][0]}\nDocument 2: {passages[1][0]}"
+	user_prompt = f"{context_str}\n\nGenerated two-hop Question from two Documents:\n"
+	messages.append(ChatMessage(role=MessageRole.USER, content=user_prompt))
+
+	completion = await client.beta.chat.completions.parse(
+		model=model_name,
+		messages=to_openai_message_dicts(messages),
+		response_format=TwoHopIncrementalResponse,
+	)
+	row["query"] = completion.choices[0].message.parsed.two_hop_question
+	return row

autorag/data/qa/query/prompt.py

@@ -0,0 +1,201 @@
+from llama_index.core.base.llms.types import ChatMessage, MessageRole
+
+QUERY_GEN_PROMPT = {
+	"factoid_single_hop": {
+		"en": [
+			ChatMessage(
+				role=MessageRole.SYSTEM,
+				content="""You're an AI tasked to convert Text into a factoid question.
+Factoid questions are those seeking brief, factual information that can be easily verified. They typically require a yes or no answer or a brief explanation and often inquire about specific details such as dates, names, places, or events.
+
+Examples of factoid questions include:
+
+- What is the capital of France?
+- Who invented the light bulb?
+- When was Wikipedia founded?
+
+Instructions:
+1. Questions MUST BE extracted from given Text
+2. Questions should be as detailed as possible from Text
+3. Create questions that ask about factual information from the Text
+4. Do not mention any of these in the questions: "in the given text", "in the provided information", etc.
+Users do not know the passage source of the question, so it should not be mentioned in the question.
+5. Do not ask about the file name or the file title. Ask about the content of the file.
+For example, avoid to write questions like `What is the file name of the document?`""",
+			)
+		],
+		"ko": [
+			ChatMessage(
+				role=MessageRole.SYSTEM,
+				content="""당신은 주어진 Text를 '사실 질문'으로 변환하는 AI입니다.
+
+사실 질문(factoid questions)이란 사실적인 정보를 요구하는 질문으로, 쉽게 검증할 수 있는 답변을 필요로 합니다. 일반적으로 예/아니오 답변이나 간단한 설명을 요구하며, 날짜, 이름, 장소 또는 사건과 같은 구체적인 세부사항에 대해 묻는 질문입니다.
+
+사실 질문의 예는 다음과 같습니다:
+
+	•	프랑스의 수도는 어디입니까?
+	•	전구를 발명한 사람은 누구입니까?
+	•	위키피디아는 언제 설립되었습니까?
+
+지침:
+	1.	질문은 반드시 주어진 Text를 기반으로 작성되어야 합니다.
+	2.	질문은 Text를 기반으로 가능한 한 구체적으로 작성되어야 합니다.
+	3.	Text에서 사실적 정보를 요구하는 질문을 만들어야 합니다. 즉, Text를 기반으로 사실 질문을 만드세요.
+	4.	질문에 “주어진 Text에서” 또는 “제공된 단락에서”와 같은 표현을 포함해서는 안 됩니다.
+사용자는 질문의 출처가 Text라는 것을 모르기 때문에 반드시 그 출처를 언급해서는 안 됩니다.
+	5.	파일 이름이나 파일 제목에 대한 질문을 하지 마세요. 파일의 내용에 대해 물어보세요.
+예를 들어, '문서의 파일 이름은 무엇입니까?'와 같은 질문을 작성하지 마세요.
+	6.	질문을 한국어로 작성하세요.""",
+			)
+		],
+		"ja": [
+			ChatMessage(
+				role=MessageRole.SYSTEM,
+				content="""あなたは与えられたTextを「実は質問」に変換するAIです。
+
+事実質問(factoid questions)とは、事実的な情報を求める質問であり、容易に検証できる回答を必要とします。 一般的に、「はい/いいえ」の返答や簡単な説明を要求し、日付、名前、場所、または事件のような具体的な詳細事項について尋ねる質問です。
+
+事実質問の例は下記の通りです:
+
+	• フランスの首都はどこですか？
+	• 電球を発明したのは誰ですか？
+	• ウィキペディアはいつ設立されましたか？
+
+指針:
+	1. 質問は、必ず与えられたTextに基づいて作成されなければなりません。
+	2. 質問は、Textに基づいて可能な限り具体的に作成されなければなりません。
+	3. Textで事実的な情報を要求する質問を作らなければなりません。 つまり、Textに基づいて事実質問を作成してください。
+	4. 質問に「与えられたTextで」または「提供された段落で」のような表現を含めてはいけません。
+ユーザーは質問の出所がTextだということを知らないので、絶対にその出所を言及してはいけません。
+	5. ファイル名やファイルのタイトルを訊かないでください。代わりに、ファイルの内容について聞いてください。
+例えば、「このドキュメントのファイル名は何ですか？」や「ファイル名は何ですか？」というような質問は書かないでください。
+	6. 質問を日本語で作成してください。""",
+			)
+		],
+	},
+	"concept_completion": {
+		"en": [
+			ChatMessage(
+				role=MessageRole.SYSTEM,
+				content="""You're an AI tasked to convert Text into a "Concept Completion" Question.
+A “concept completion” question asks directly about the essence or identity of a concept.
+
+Follow the following instructions.
+Instructions:
+1. Questions MUST BE extracted from given Text
+2. Questions should be as detailed as possible from Text
+3. Create questions that ask about information from the Text
+4. MUST include specific keywords from the Text.
+5. Do not mention any of these in the questions: "in the given text", "in the provided information", etc.
+Users do not know the passage source of the question, so it should not be mentioned in the question.
+6. Do not ask about the file name or the file title. Ask about the content of the file.
+For example, avoid to write questions like `What is the file name of the document?""",
+			)
+		],
+		"ko": [
+			ChatMessage(
+				role=MessageRole.SYSTEM,
+				content="""당신은 Text를 “개념 완성” 질문으로 변환하는 AI입니다.
+"개념 완성" 질문은 개념의 본질이나 정체성에 대해 직접적으로 묻는 질문입니다.
+
+다음 지시사항을 따르세요.
+지시사항:
+1.	질문은 반드시 주어진 Text를 기반으로 작성되어야 합니다.
+2.	질문은 Text를 기반으로 가능한 한 자세하게 작성되어야 합니다.
+3.	Text에서 제공된 정보를 묻는 질문을 생성하세요.
+4.	Text의 특정 키워드를 반드시 질문에 포함하세요.
+5.	질문에 “주어진 Text에서” 또는 “제공된 단락에서”와 같은 표현을 포함해서는 안 됩니다.
+사용자는 질문의 출처가 Text라는 것을 모르기 때문에 반드시 그 출처를 언급해서는 안 됩니다.
+6.	파일 이름이나 파일 제목에 대한 질문을 하지 마세요. 파일의 내용에 대해 물어보세요.
+예를 들어, '문서의 파일 이름은 무엇입니까?'와 같은 질문을 작성하지 마세요.
+7.	질문을 한국어로 작성하세요.""",
+			)
+		],
+		"ja": [
+			ChatMessage(
+				role=MessageRole.SYSTEM,
+				content="""あなたはTextを「概念完成」の質問に変換するAIです。
+「概念完成」の質問は概念の本質やアイデンティティについて直接に尋ねる質問です。
+
+次の指示に従ってください。
+指示事項:
+1. 質問は、必ず与えられたTextに基づいて作成されなければなりません。
+2. 質問は、Textに基づいてできるだけ詳しく作成されなければなりません。
+3. Textで提供された情報を尋ねる質問を作成してください。
+4. Textの特定のキーワードを必ず質問に含めてください。
+5. 質問に「与えられたTextで」または「提供された段落で」のような表現を含めてはいけません。
+ユーザーは質問の出所がTextだということを知らないので、絶対にその出所を言及してはいけません。
+6. ファイル名やファイルのタイトルを訊かないでください。代わりに、ファイルの内容について聞いてください。
+例えば、「このドキュメントのファイル名は何ですか？」や「ファイル名は何ですか？」というような質問は書かないでください。
+7. 質問を日本語で作成してください。""",
+			)
+		],
+	},
+	"two_hop_incremental": {
+		"en": [
+			ChatMessage(
+				role=MessageRole.SYSTEM,
+				content="Generate a multi-hop question for the given answer which requires reference to all of the given documents.",
+			),
+			ChatMessage(
+				role=MessageRole.USER,
+				content="""Document 1: The Municipality of Nuevo Laredo is located in the Mexican state of Tamaulipas.
+Document 2: The Ciudad Deportiva (Sports City ¨ ¨) is a sports
+complex in Nuevo Laredo, Mexico. It is home to the Tecolotes de
+Nuevo Laredo Mexican Baseball League team and ...""",
+			),
+			ChatMessage(
+				role=MessageRole.ASSISTANT,
+				content="""Answer: Tamaulipas
+One-hop question (using Document 1): In which Mexican state is Nuevo Laredo located?
+Two-hop question (using Document 2):  In which Mexican state can one find the Ciudad Deportiva, home to the Tecolotes de Nuevo Laredo?""",
+			),
+		],
+		"ko": [
+			ChatMessage(
+				role=MessageRole.SYSTEM,
+				content="Generate a multi-hop question for the given answer which requires reference to all of the given documents.",
+			),
+			ChatMessage(
+				role=MessageRole.USER,
+				content="""Document 1: The Municipality of Nuevo Laredo is located in the Mexican state of Tamaulipas.
+Document 2: The Ciudad Deportiva (Sports City ¨ ¨) is a sports
+complex in Nuevo Laredo, Mexico. It is home to the Tecolotes de
+Nuevo Laredo Mexican Baseball League team and ...""",
+			),
+			ChatMessage(
+				role=MessageRole.ASSISTANT,
+				content="""Answer: Tamaulipas
+One-hop question (using Document 1): In which Mexican state is Nuevo Laredo located?
+Two-hop question (using Document 2):  In which Mexican state can one find the Ciudad Deportiva, home to the Tecolotes de Nuevo Laredo?""",
+			),
+		],
+		"ja": [
+			ChatMessage(
+				role=MessageRole.SYSTEM,
+				content="与えられた答えに対するマルチホップ質問を生成し、与えられたすべての文書を参照する必要があります。",
+			),
+			ChatMessage(
+				role=MessageRole.USER,
+				content="""Document 1: ヌエヴォ·ラレド自治体はメキシコのタマウリパス州にあります。
+シウダー・デポルティーバ（スポーツ・シティ）は、
+メキシコのヌエボ・ラレドにある複合スポーツ施設です。""",
+			),
+			ChatMessage(
+				role=MessageRole.ASSISTANT,
+				content="""Answer: Tamaulipas
+One-hop question (using Document 1): ヌエヴォ·ラレド自治体はどのメキシコの州にありますか？
+Two-hop question (using Document 2): ヌエヴォ·ラレドのテコロテス·デ·テコロテスの故郷であるメキシコの州はどこですか？""",
+			),
+		],
+	},
+}
+
+# Experimental feature
+QUERY_GEN_PROMPT_EXTRA = {
+	"multiple_queries": {
+		"en": "\nAdditional instructions:\n  - Please make {n} questions.",
+		"ko": "\n추가 지침:\n  - 질문은 {n}개를 만드세요.",
+		"ja": "\n追加指示:\n  - 質問を{n}個作成してください。",
+	}
+}

autorag/data/qa/sample.py

@@ -0,0 +1,26 @@
+import uuid
+from typing import Iterable
+
+import pandas as pd
+
+
+def random_single_hop(
+	corpus_df: pd.DataFrame, n: int, random_state: int = 42
+) -> pd.DataFrame:
+	sample_df = corpus_df.sample(n, random_state=random_state)
+	return pd.DataFrame(
+		{
+			"qid": [str(uuid.uuid4()) for _ in range(len(sample_df))],
+			"retrieval_gt": [[[id_]] for id_ in sample_df["doc_id"].tolist()],
+		}
+	)
+
+
+def range_single_hop(corpus_df: pd.DataFrame, idx_range: Iterable):
+	sample_df = corpus_df.iloc[idx_range]
+	return pd.DataFrame(
+		{
+			"qid": [str(uuid.uuid4()) for _ in range(len(sample_df))],
+			"retrieval_gt": [[[id_]] for id_ in sample_df["doc_id"].tolist()],
+		}
+	)