thinhlpg
4 weeks ago
parent
5eabd121a3
commit
89e07bc02d
@ -1,2 +1,3 @@
|
||||
HF_TOKEN=<your-huggingface-token>
|
||||
OPENROUTER_API_KEY=<your-openrouter-api-key>
|
||||
TAVILY_API_KEY=<your-tavily-api-key>
|
||||
SERPER_API_KEY=<your-serper-api-key>
|
||||
@ -1,3 +0,0 @@
|
||||
[submodule "third_party/FlashRAG"]
|
||||
path = third_party/FlashRAG
|
||||
url = https://github.com/RUC-NLPIR/FlashRAG.git
|
||||
@ -1,656 +0,0 @@
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from copy import deepcopy
|
||||
from datetime import datetime
|
||||
from functools import wraps
|
||||
|
||||
import numpy as np
|
||||
import requests
|
||||
from flashrag.config import Config
|
||||
from flashrag.generator.generator import BaseGenerator
|
||||
from flashrag.pipeline import BasicPipeline
|
||||
from flashrag.retriever.retriever import BaseTextRetriever
|
||||
from flashrag.utils import get_dataset
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from config import logger
|
||||
from src.agent import Agent, AgenticOutputs
|
||||
from src.tokenizer_adapter import LlamaTokenizerAdapter, R1DistilTokenizerAdapter
|
||||
|
||||
|
||||
def retry(max_retries=10, sleep=1):
|
||||
"""Decorator to retry a function with exponential backoff."""
|
||||
|
||||
def decorator(func):
|
||||
@wraps(func)
|
||||
def wrapper(*args, **kwargs):
|
||||
func_name = func.__name__
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
result = func(*args, **kwargs)
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.warning(f"Attempt {attempt + 1} of {func_name} failed: {e}")
|
||||
if attempt == max_retries - 1:
|
||||
logger.error(f"Function {func_name} failed after {max_retries} retries.", exc_info=True)
|
||||
raise e
|
||||
backoff_time = sleep * (2**attempt)
|
||||
logger.info(f"Retrying {func_name} in {backoff_time:.2f} seconds...")
|
||||
time.sleep(backoff_time)
|
||||
logger.error(f"Function {func_name} retry logic finished unexpectedly.")
|
||||
return None
|
||||
|
||||
return wrapper
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
class RemoteRetriever(BaseTextRetriever):
|
||||
"""A wrapper for remote retriever service with retry logic and logging."""
|
||||
|
||||
def __init__(self, config: Config):
|
||||
"""Initializes the RemoteRetriever."""
|
||||
super().__init__(config)
|
||||
self.remote_url = f"http://{getattr(config, 'remote_retriever_url', 'localhost:8001')}"
|
||||
self.topk = getattr(config, "retriever_topk", 5)
|
||||
logger.info(f"🔗 Remote retriever URL: {self.remote_url}")
|
||||
|
||||
@retry(max_retries=3, sleep=2)
|
||||
def _search(self, query: str, num: int | None = None, return_score: bool = False) -> list[dict]:
|
||||
"""Search for documents using the remote retriever service."""
|
||||
num = num if num is not None else self.topk
|
||||
url = f"{self.remote_url}/search"
|
||||
|
||||
try:
|
||||
response = requests.post(
|
||||
url,
|
||||
json={"query": query, "top_n": num, "return_score": return_score},
|
||||
timeout=30,
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
results = response.json()
|
||||
return results
|
||||
except requests.exceptions.Timeout:
|
||||
logger.error(f"Search request timed out after 30 seconds for query: {query[:50]}...")
|
||||
raise
|
||||
except requests.exceptions.ConnectionError:
|
||||
logger.error(f"Could not connect to search service at {url}")
|
||||
raise
|
||||
except requests.exceptions.RequestException as e:
|
||||
logger.error(f"Search request failed: {e}", exc_info=True)
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Unexpected search error: {str(e)}", exc_info=True)
|
||||
raise
|
||||
|
||||
@retry(max_retries=3, sleep=2)
|
||||
def _batch_search(
|
||||
self, queries: list[str], num: int | None = None, return_score: bool = False
|
||||
) -> list[list[dict]]:
|
||||
"""Batch search for documents using the remote retriever service."""
|
||||
num = num if num is not None else self.topk
|
||||
url = f"{self.remote_url}/batch_search"
|
||||
|
||||
try:
|
||||
response = requests.post(
|
||||
url,
|
||||
json={"query": queries, "top_n": num, "return_score": return_score},
|
||||
timeout=60,
|
||||
)
|
||||
response.raise_for_status()
|
||||
results = response.json()
|
||||
return results
|
||||
except requests.exceptions.Timeout:
|
||||
logger.error(f"Batch search request timed out after 60 seconds for {len(queries)} queries.")
|
||||
raise
|
||||
except requests.exceptions.ConnectionError:
|
||||
logger.error(f"Could not connect to batch search service at {url}")
|
||||
raise
|
||||
except requests.exceptions.RequestException as e:
|
||||
logger.error(f"Batch search request failed: {e}", exc_info=True)
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Unexpected batch search error: {str(e)}", exc_info=True)
|
||||
raise
|
||||
|
||||
|
||||
class ReSearchPipeline(BasicPipeline):
|
||||
"""Pipeline for ReSearch method using Agent for generation and tool use."""
|
||||
|
||||
def __init__(
|
||||
self, config: Config, retriever: BaseTextRetriever | None = None, generator: BaseGenerator | None = None
|
||||
):
|
||||
"""Initializes the ReSearchPipeline."""
|
||||
super().__init__(config)
|
||||
logger.info("🔧 Initializing ReSearchPipeline...")
|
||||
|
||||
self.retriever = retriever or RemoteRetriever(config)
|
||||
|
||||
self.generator = generator or SGLRemoteGenerator(config)
|
||||
|
||||
try:
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(config.generator_model_path, trust_remote_code=True)
|
||||
if not self.tokenizer.pad_token:
|
||||
logger.warning("Tokenizer does not have a pad token; setting to eos_token.")
|
||||
self.tokenizer.pad_token = self.tokenizer.eos_token
|
||||
self.tokenizer.padding_side = "left"
|
||||
logger.info("✅ Tokenizer initialized.")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to initialize tokenizer: {e}", exc_info=True)
|
||||
raise
|
||||
|
||||
tokenizer_name = self.tokenizer.name_or_path.lower()
|
||||
|
||||
if "deepseek-ai/deepseek-r1-distill" in tokenizer_name:
|
||||
adapter = R1DistilTokenizerAdapter()
|
||||
elif "llama" in tokenizer_name:
|
||||
adapter = LlamaTokenizerAdapter()
|
||||
else:
|
||||
logger.warning(f"Unknown tokenizer type '{tokenizer_name}', defaulting to R1DistilTokenizerAdapter.")
|
||||
adapter = R1DistilTokenizerAdapter()
|
||||
logger.info(f"🔩 Using Tokenizer Adapter: {type(adapter).__name__}")
|
||||
|
||||
def retriever_search(query: str, return_type=str, results: int = 5):
|
||||
try:
|
||||
search_results = self.retriever._search(query, num=results)
|
||||
return self.format_search_results(search_results)
|
||||
except Exception as e:
|
||||
logger.error(f"Error during agent's retriever search for query '{query[:50]}...': {e}", exc_info=True)
|
||||
return "<information>Search failed due to an internal error."
|
||||
|
||||
self.agent = Agent(adapter, search_fn=retriever_search)
|
||||
logger.info("✅ Agent initialized.")
|
||||
logger.info("✅ ReSearchPipeline initialized successfully.")
|
||||
|
||||
def format_search_results(self, search_results: list[dict]) -> str:
|
||||
"""Formats search results into a string for the agent prompt."""
|
||||
if not search_results:
|
||||
return "<information>No results found.</information>"
|
||||
max_content_len = 500
|
||||
formatted = "\n-------\n".join(
|
||||
[
|
||||
f"Result {i + 1}: {r.get('contents', 'N/A')[:max_content_len]}{'...' if len(r.get('contents', '')) > max_content_len else ''}"
|
||||
for i, r in enumerate(search_results)
|
||||
]
|
||||
)
|
||||
formatted_str = f"<information>{formatted}</information>"
|
||||
|
||||
return formatted_str
|
||||
|
||||
def extract_search_query(self, text: str) -> str | None:
|
||||
"""Extract search query from text between <search> tags."""
|
||||
pattern = re.compile(r"<search>(.*?)</search>", re.DOTALL)
|
||||
matches = pattern.findall(text)
|
||||
if matches:
|
||||
query = matches[-1].strip()
|
||||
return query
|
||||
return None
|
||||
|
||||
def extract_answer(self, text: str) -> str | None:
|
||||
"""Extract answer from text between <answer> tags."""
|
||||
pattern = re.compile(r"<answer>(.*?)</answer>", re.DOTALL)
|
||||
matches = pattern.findall(text)
|
||||
if matches:
|
||||
answer = matches[-1].strip()
|
||||
|
||||
return answer
|
||||
|
||||
return None
|
||||
|
||||
def run(self, dataset, do_eval: bool = True, pred_process_fun=None):
|
||||
"""Runs the ReSearch pipeline on the dataset using the Agent."""
|
||||
logger.info(f"🏃 Starting ReSearch pipeline run with {len(dataset)} items...")
|
||||
|
||||
try:
|
||||
questions = [item.question if hasattr(item, "question") else item["question"] for item in dataset]
|
||||
|
||||
except (KeyError, AttributeError, TypeError) as e:
|
||||
logger.error(f"Failed to extract questions from dataset items. Error: {e}", exc_info=True)
|
||||
logger.error("Ensure dataset items have a 'question' key or attribute.")
|
||||
return dataset
|
||||
|
||||
agent_max_generations = getattr(self.config, "agent_max_generations", 32)
|
||||
generator_max_output_len = getattr(self.config, "generator_max_output_len", 24576)
|
||||
|
||||
try:
|
||||
logger.info(f"🤖 Running agent inference for {len(questions)} questions...")
|
||||
agent_outputs: AgenticOutputs = self.agent.run_agent(
|
||||
generate_fn=self.generator.generate,
|
||||
tokenizer=self.tokenizer,
|
||||
questions=questions,
|
||||
max_generations=agent_max_generations,
|
||||
max_new_tokens=generator_max_output_len,
|
||||
)
|
||||
final_responses = agent_outputs.final_response_str
|
||||
logger.info(f"✅ Agent inference completed. Received {len(final_responses)} final responses.")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Agent run failed during inference: {e}", exc_info=True)
|
||||
logger.warning("Agent run failed, attempting evaluation with potentially incomplete results.")
|
||||
for item in dataset:
|
||||
if hasattr(item, "update_output"):
|
||||
item.update_output("pred", "AGENT_ERROR")
|
||||
elif isinstance(item, dict):
|
||||
item["pred"] = "AGENT_ERROR"
|
||||
|
||||
logger.info("📝 Extracting answers and updating dataset items...")
|
||||
num_updated = 0
|
||||
num_missing_answers = 0
|
||||
if len(final_responses) == len(dataset):
|
||||
for i, item in enumerate(dataset):
|
||||
response = final_responses[i]
|
||||
answer = self.extract_answer(response)
|
||||
pred_to_save = answer if answer is not None else ""
|
||||
|
||||
if answer is None:
|
||||
num_missing_answers += 1
|
||||
|
||||
if hasattr(item, "update_output"):
|
||||
item.update_output("pred", pred_to_save)
|
||||
item.update_output("final_response", response)
|
||||
num_updated += 1
|
||||
elif isinstance(item, dict):
|
||||
item["pred"] = pred_to_save
|
||||
item["final_response"] = response
|
||||
num_updated += 1
|
||||
else:
|
||||
logger.warning(f"Item {i} has unknown type {type(item)}, cannot update with prediction.")
|
||||
|
||||
logger.info(f"Updated {num_updated}/{len(dataset)} dataset items with predictions.")
|
||||
if num_missing_answers > 0:
|
||||
logger.warning(f"{num_missing_answers} items had no <answer> tag.")
|
||||
else:
|
||||
logger.error(
|
||||
f"Mismatch between dataset size ({len(dataset)}) and number of agent responses ({len(final_responses)}). Cannot reliably update dataset."
|
||||
)
|
||||
for item in dataset:
|
||||
if hasattr(item, "update_output"):
|
||||
item.update_output("pred", "RESPONSE_COUNT_MISMATCH")
|
||||
elif isinstance(item, dict):
|
||||
item["pred"] = "RESPONSE_COUNT_MISMATCH"
|
||||
|
||||
if do_eval:
|
||||
logger.info("📊 Evaluating results using BasicPipeline.evaluate...")
|
||||
try:
|
||||
dataset = self.evaluate(dataset, do_eval=True, pred_process_fun=pred_process_fun)
|
||||
logger.info("✅ Evaluation completed via base class method.")
|
||||
except Exception as e:
|
||||
logger.error(f"Error during BasicPipeline.evaluate: {e}", exc_info=True)
|
||||
logger.warning("Evaluation may be incomplete.")
|
||||
else:
|
||||
logger.info("Skipping evaluation step as do_eval=False.")
|
||||
|
||||
logger.info("✅ ReSearch pipeline run finished.")
|
||||
return dataset
|
||||
|
||||
|
||||
class SGLRemoteGenerator(BaseGenerator):
|
||||
"""Class for decoder-only generator, based on SGLang remote service."""
|
||||
|
||||
def __init__(self, config: Config):
|
||||
"""Initializes the SGLRemoteGenerator."""
|
||||
super().__init__(config)
|
||||
logger.info("🔧 Initializing SGLRemoteGenerator...")
|
||||
sgl_url = getattr(config, "sgl_remote_url", "localhost:8002")
|
||||
self.sgl_remote_url = f"http://{sgl_url}/generate"
|
||||
self.health_check_url = f"http://{sgl_url}/health"
|
||||
logger.info(f"🔗 Remote Generator URL: {self.sgl_remote_url}")
|
||||
self.model_path = getattr(config, "generator_model_path", None)
|
||||
if not self.model_path:
|
||||
logger.error("generator_model_path not found in config!")
|
||||
raise ValueError("generator_model_path is required for SGLRemoteGenerator")
|
||||
|
||||
try:
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(self.model_path, trust_remote_code=True)
|
||||
logger.info("✅ Tokenizer loaded for generator.")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load tokenizer for generator from {self.model_path}: {e}", exc_info=True)
|
||||
raise
|
||||
|
||||
self.generation_params = getattr(config, "generation_params", {})
|
||||
self.config = config
|
||||
|
||||
self._check_health()
|
||||
|
||||
def _check_health(self):
|
||||
"""Checks the health of the remote generator service."""
|
||||
try:
|
||||
test_response = requests.get(self.health_check_url, timeout=5)
|
||||
test_response.raise_for_status()
|
||||
logger.info("✅ Remote generator service is available")
|
||||
except requests.exceptions.RequestException as e:
|
||||
logger.error(f"Could not connect or verify remote generator service at {self.health_check_url}: {str(e)}")
|
||||
logger.warning("Please ensure the SGLang service is running and accessible.")
|
||||
|
||||
@retry(max_retries=5, sleep=2)
|
||||
def generate(
|
||||
self,
|
||||
input_list: list[str] | str,
|
||||
return_raw_output: bool = False,
|
||||
return_scores: bool = False,
|
||||
**params,
|
||||
) -> list[str] | tuple[list[str], list[list[float]]] | list[dict]:
|
||||
"""Generates text using the remote SGLang service."""
|
||||
if isinstance(input_list, str):
|
||||
input_list = [input_list]
|
||||
if not isinstance(input_list, list) or not all(isinstance(item, str) for item in input_list):
|
||||
raise ValueError("Input must be a string or a list of strings.")
|
||||
|
||||
batch_size = len(input_list)
|
||||
data_to_remote = {"text": input_list}
|
||||
|
||||
effective_params = deepcopy(self.generation_params)
|
||||
effective_params.update(params)
|
||||
|
||||
curr_sampling_params = {}
|
||||
if effective_params.get("do_sample", True) is False:
|
||||
curr_sampling_params["temperature"] = 0.0
|
||||
else:
|
||||
curr_sampling_params["temperature"] = effective_params.get(
|
||||
"temperature", getattr(self.config, "temperature", 0.7)
|
||||
)
|
||||
|
||||
default_max_tokens = getattr(self.config, "generator_max_output_len", 1024)
|
||||
curr_sampling_params["max_new_tokens"] = effective_params.get("max_new_tokens", default_max_tokens)
|
||||
|
||||
stop_sequences = effective_params.get("stop", [])
|
||||
if isinstance(stop_sequences, str):
|
||||
stop_sequences = [stop_sequences]
|
||||
if stop_sequences:
|
||||
curr_sampling_params["stop"] = stop_sequences
|
||||
|
||||
keys_to_remove = ["do_sample", "temperature", "max_new_tokens", "stop"]
|
||||
for key in keys_to_remove:
|
||||
effective_params.pop(key, None)
|
||||
|
||||
if "top_p" in effective_params:
|
||||
curr_sampling_params["top_p"] = effective_params["top_p"]
|
||||
if "top_k" in effective_params:
|
||||
curr_sampling_params["top_k"] = effective_params["top_k"]
|
||||
|
||||
data_to_remote["sampling_params"] = curr_sampling_params
|
||||
|
||||
if return_scores:
|
||||
data_to_remote["return_logprob"] = True
|
||||
data_to_remote["top_logprobs_num"] = getattr(self.config, "top_logprobs_num", 2)
|
||||
|
||||
try:
|
||||
response = requests.post(
|
||||
self.sgl_remote_url, json=data_to_remote, timeout=120, headers={"Content-Type": "application/json"}
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
response_list = response.json()
|
||||
|
||||
if return_raw_output:
|
||||
return response_list
|
||||
|
||||
generated_text = []
|
||||
for item in response_list:
|
||||
text = item.get("text", "")
|
||||
finish_reason = item.get("meta_info", {}).get("finish_reason", {})
|
||||
matched_stop = finish_reason.get("matched")
|
||||
if matched_stop and curr_sampling_params.get("stop") and matched_stop in curr_sampling_params["stop"]:
|
||||
text += matched_stop
|
||||
generated_text.append(text)
|
||||
|
||||
if return_scores:
|
||||
scores = []
|
||||
for resp_item in response_list:
|
||||
logprobs_list = resp_item.get("meta_info", {}).get("output_token_logprobs", [])
|
||||
token_scores = [
|
||||
np.exp(logprob[0]) if (logprob and len(logprob) > 0) else 0.0 for logprob in logprobs_list
|
||||
]
|
||||
scores.append(token_scores)
|
||||
return generated_text, scores
|
||||
else:
|
||||
return generated_text
|
||||
|
||||
except requests.exceptions.Timeout:
|
||||
logger.error("Generation request timed out after 120 seconds.")
|
||||
raise
|
||||
except requests.exceptions.ConnectionError:
|
||||
logger.error(f"Could not connect to remote generator service at {self.sgl_remote_url}.")
|
||||
raise
|
||||
except requests.exceptions.RequestException as e:
|
||||
logger.error(f"Network error during generation: {str(e)}", exc_info=True)
|
||||
raise
|
||||
except json.JSONDecodeError:
|
||||
response_text = "Unknown (error occurred before response object assignment)"
|
||||
if "response" in locals() and hasattr(response, "text"):
|
||||
response_text = response.text[:500]
|
||||
logger.error(
|
||||
f"Failed to decode JSON response from {self.sgl_remote_url}. Response text: {response_text}...",
|
||||
exc_info=True,
|
||||
)
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Unexpected error during generation: {str(e)}", exc_info=True)
|
||||
raise
|
||||
|
||||
|
||||
def load_dataset_items(config: Config, split: str) -> list[dict | object]:
|
||||
"""Loads dataset items using flashrag's get_dataset."""
|
||||
logger.info(f"📚 Loading dataset: {config.dataset_name}, Split: {split}")
|
||||
try:
|
||||
all_splits = get_dataset(config)
|
||||
if split not in all_splits:
|
||||
logger.error(
|
||||
f"Split '{split}' not found in dataset '{config.dataset_name}'. Available splits: {list(all_splits.keys())}"
|
||||
)
|
||||
return []
|
||||
dataset_items = all_splits[split]
|
||||
logger.info(f"Successfully loaded {len(dataset_items)} items for split '{split}'.")
|
||||
|
||||
return dataset_items
|
||||
except FileNotFoundError:
|
||||
logger.error(
|
||||
f"Dataset files not found for '{config.dataset_name}' in '{config.data_dir}'. Check config and paths."
|
||||
)
|
||||
return []
|
||||
except Exception as e:
|
||||
logger.error(f"Error loading dataset using get_dataset: {e}", exc_info=True)
|
||||
return []
|
||||
|
||||
|
||||
def save_results(args: argparse.Namespace, config: Config, result_dataset, run_duration: float):
|
||||
"""Saves summary and debug information."""
|
||||
logger.info("💾 Saving results...")
|
||||
summary_file = os.path.join(args.save_dir, f"{args.save_note}_summary.txt")
|
||||
debug_file = os.path.join(args.save_dir, f"{args.save_note}_debug.json")
|
||||
|
||||
num_items = len(result_dataset)
|
||||
|
||||
logger.info(f"Saving summary results to {summary_file}...")
|
||||
try:
|
||||
with open(summary_file, "w", encoding="utf-8") as f:
|
||||
f.write("EVALUATION SUMMARY\n")
|
||||
f.write("=================\n\n")
|
||||
f.write(f"Timestamp: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
|
||||
f.write(f"Run Duration: {run_duration:.2f} seconds\n")
|
||||
f.write(f"Dataset: {config.dataset_name} ({args.split} split)\n")
|
||||
f.write(f"Model: {config.generator_model_path}\n")
|
||||
f.write(f"Retriever: {config.remote_retriever_url}\n")
|
||||
f.write(f"Agent Max Generations: {getattr(config, 'agent_max_generations', 'N/A')}\n")
|
||||
f.write(f"Generator Max Output Len: {getattr(config, 'generator_max_output_len', 'N/A')}\n\n")
|
||||
f.write(f"Total items processed: {num_items}\n")
|
||||
f.write("\nNote: Verification was skipped in this run.\n")
|
||||
f.write("Note: Overall metrics (like EM, F1) are usually printed to console by evaluate method.\n")
|
||||
|
||||
logger.info(f"✅ Summary saved to {summary_file}")
|
||||
except Exception as e:
|
||||
logger.error(f"Error saving summary file '{summary_file}': {e}", exc_info=True)
|
||||
|
||||
logger.info(f"Saving debug information (predictions & responses) to {debug_file}...")
|
||||
try:
|
||||
debug_data = []
|
||||
for i, item in enumerate(result_dataset):
|
||||
item_data: dict[str, object] = {}
|
||||
|
||||
def get_item_value(data_item, key_or_attr: str) -> str | int | float | list | bool | None:
|
||||
if isinstance(data_item, dict):
|
||||
return data_item.get(key_or_attr)
|
||||
elif hasattr(data_item, key_or_attr):
|
||||
return getattr(data_item, key_or_attr)
|
||||
return None
|
||||
|
||||
item_data["item_index"] = i
|
||||
item_data["question"] = get_item_value(item, "question")
|
||||
item_data["prediction"] = get_item_value(item, "pred")
|
||||
item_data["final_response"] = get_item_value(item, "final_response")
|
||||
|
||||
gt_answer_val = None
|
||||
try:
|
||||
gt_answer_val = get_item_value(item, "answer")
|
||||
if gt_answer_val is None:
|
||||
answers_list = get_item_value(item, "answers")
|
||||
if isinstance(answers_list, list) and answers_list:
|
||||
raw_ans = answers_list[0]
|
||||
if isinstance(raw_ans, (str, int, float, bool)):
|
||||
gt_answer_val = raw_ans
|
||||
else:
|
||||
gt_answer_val = str(raw_ans)
|
||||
elif not isinstance(gt_answer_val, (str, int, float, bool)):
|
||||
gt_answer_val = str(gt_answer_val)
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not safely get ground truth for item {i}: {e}")
|
||||
gt_answer_val = "ERROR_GETTING_ANSWER"
|
||||
item_data["ground_truth"] = gt_answer_val
|
||||
|
||||
eval_score_val = None
|
||||
try:
|
||||
eval_score_val = get_item_value(item, "score")
|
||||
if not isinstance(eval_score_val, (str, int, float, bool, type(None))):
|
||||
eval_score_val = str(eval_score_val)
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not safely get score for item {i}: {e}")
|
||||
eval_score_val = "ERROR_GETTING_SCORE"
|
||||
item_data["eval_score"] = eval_score_val
|
||||
|
||||
debug_data.append(item_data)
|
||||
|
||||
with open(debug_file, "w", encoding="utf-8") as f:
|
||||
json.dump(debug_data, f, indent=2, ensure_ascii=False)
|
||||
logger.info(f"✅ Debug information saved to {debug_file}")
|
||||
except Exception as e:
|
||||
logger.error(f"Error saving debug file '{debug_file}': {e}", exc_info=True)
|
||||
|
||||
|
||||
def research(args: argparse.Namespace, config: Config):
|
||||
"""Main function to run the research evaluation pipeline."""
|
||||
logger.info("🚀 Starting research pipeline execution...")
|
||||
start_time = time.time()
|
||||
|
||||
test_data = load_dataset_items(config, args.split)
|
||||
if not test_data:
|
||||
logger.error("Failed to load test data. Exiting.")
|
||||
return
|
||||
|
||||
try:
|
||||
logger.info("🏗️ Building ReSearchPipeline...")
|
||||
pipeline = ReSearchPipeline(config)
|
||||
logger.info("✅ Pipeline built successfully.")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to initialize ReSearchPipeline: {e}", exc_info=True)
|
||||
return
|
||||
|
||||
agent_max_generations = getattr(config, "agent_max_generations", 32)
|
||||
generator_max_output_len = getattr(config, "generator_max_output_len", 24576)
|
||||
|
||||
try:
|
||||
logger.info("🏃 Starting pipeline run...")
|
||||
result_dataset = pipeline.run(test_data, do_eval=True)
|
||||
logger.info("✅ Pipeline run completed.")
|
||||
except Exception as e:
|
||||
logger.error(f"Error during pipeline run: {e}", exc_info=True)
|
||||
result_dataset = test_data
|
||||
logger.warning("Pipeline run failed, attempting to save inputs/partial results.")
|
||||
|
||||
run_duration = time.time() - start_time
|
||||
logger.info(f"Total run duration: {run_duration:.2f} seconds.")
|
||||
save_results(args, config, result_dataset, run_duration)
|
||||
|
||||
logger.info("🏁 Research pipeline execution finished.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Running ReSearch Evaluation Pipeline")
|
||||
parser.add_argument(
|
||||
"--config_path", type=str, default="./eval_config.yaml", help="Path to the main FlashRAG config file."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dataset_name",
|
||||
type=str,
|
||||
default="bamboogle",
|
||||
help="Name of the dataset (must match config or data_dir structure).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--split", type=str, default="test", help="Dataset split to evaluate (e.g., test, validation)."
|
||||
)
|
||||
parser.add_argument("--save_dir", type=str, default="./output_logs", help="Directory to save logs and results.")
|
||||
parser.add_argument("--save_note", type=str, default="research_run", help="A note to prepend to saved filenames.")
|
||||
|
||||
parser.add_argument("--data_dir", type=str, help="Override data directory specified in config.")
|
||||
parser.add_argument(
|
||||
"--sgl_remote_url", type=str, help="Override SGLang remote generator URL (e.g., localhost:8002)."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--remote_retriever_url", type=str, help="Override remote retriever URL (e.g., localhost:8001)."
|
||||
)
|
||||
parser.add_argument("--generator_model_path", type=str, help="Override generator model path specified in config.")
|
||||
parser.add_argument("--retriever_topk", type=int, help="Override retriever top K.")
|
||||
parser.add_argument("--generator_max_output_len", type=int, help="Override generator max output length.")
|
||||
parser.add_argument("--agent_max_generations", type=int, help="Override agent max interaction turns.")
|
||||
|
||||
args = parser.parse_args()
|
||||
logger.info(f"Starting evaluation script with arguments: {args}")
|
||||
|
||||
try:
|
||||
os.makedirs(args.save_dir, exist_ok=True)
|
||||
logger.info(f"💾 Logs and results will be saved to: {args.save_dir}")
|
||||
except OSError as e:
|
||||
logger.error(f"Could not create save directory '{args.save_dir}': {e}", exc_info=True)
|
||||
exit(1)
|
||||
|
||||
config_overrides = {
|
||||
k: v
|
||||
for k, v in vars(args).items()
|
||||
if v is not None
|
||||
and k
|
||||
not in [
|
||||
"config_path",
|
||||
"dataset_name",
|
||||
"split",
|
||||
"save_dir",
|
||||
"save_note",
|
||||
]
|
||||
}
|
||||
|
||||
logger.info(f"🔧 Loading configuration from: {args.config_path}")
|
||||
try:
|
||||
config = Config(args.config_path, config_dict=config_overrides)
|
||||
config.dataset_name = args.dataset_name
|
||||
if args.data_dir:
|
||||
config.data_dir = args.data_dir
|
||||
|
||||
logger.info(f"Effective data_dir: {getattr(config, 'data_dir', 'N/A')}")
|
||||
logger.info(f"Effective generator_model_path: {getattr(config, 'generator_model_path', 'N/A')}")
|
||||
logger.info(f"Effective sgl_remote_url: {getattr(config, 'sgl_remote_url', 'N/A')}")
|
||||
logger.info(f"Effective remote_retriever_url: {getattr(config, 'remote_retriever_url', 'N/A')}")
|
||||
|
||||
logger.info("✅ Config loaded and potentially overridden by CLI arguments.")
|
||||
|
||||
config["dataset_path"] = os.path.join(config.data_dir, config.dataset_name)
|
||||
|
||||
except FileNotFoundError:
|
||||
logger.error(f"Config file not found at '{args.config_path}'. Please check the path.")
|
||||
exit(1)
|
||||
except Exception as e:
|
||||
logger.error(f"Error loading or processing configuration: {e}", exc_info=True)
|
||||
exit(1)
|
||||
|
||||
research(args, config)
|
||||
@ -1,68 +0,0 @@
|
||||
import argparse
|
||||
import os
|
||||
import zipfile
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
from config import DATA_DIR
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
"""Parse command line arguments.
|
||||
|
||||
Returns:
|
||||
argparse.Namespace: Parsed arguments
|
||||
"""
|
||||
parser = argparse.ArgumentParser(description="Download FlashRAG datasets from HuggingFace Hub")
|
||||
parser.add_argument(
|
||||
"--repo-id",
|
||||
type=str,
|
||||
default="RUC-NLPIR/FlashRAG_datasets",
|
||||
help="HuggingFace repository IDs",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--local-dir",
|
||||
type=str,
|
||||
default=DATA_DIR / "flashrag_datasets",
|
||||
help="Local directory to save model",
|
||||
)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
"""Main function to download model."""
|
||||
args = parse_args()
|
||||
load_dotenv(override=True)
|
||||
|
||||
# Configuration
|
||||
HF_TOKEN = os.getenv("HF_TOKEN")
|
||||
|
||||
ALLOW_PATTERNS = [
|
||||
"*retrieval-corpus*",
|
||||
"*bamboogle*",
|
||||
"*nq*",
|
||||
]
|
||||
|
||||
# Download the model
|
||||
snapshot_download(
|
||||
token=HF_TOKEN,
|
||||
repo_id=args.repo_id,
|
||||
local_dir=args.local_dir,
|
||||
repo_type="dataset",
|
||||
# ignore_patterns=IGNORE_PATTERNS,
|
||||
allow_patterns=ALLOW_PATTERNS,
|
||||
)
|
||||
|
||||
# unzip data/flashrag_datasets/retrieval-corpus/wiki18_100w.zip
|
||||
print("Unzipping wiki18_100w.zip. Might take a while...")
|
||||
zip_file_path = os.path.join(args.local_dir, "retrieval-corpus", "wiki18_100w.zip")
|
||||
with zipfile.ZipFile(zip_file_path, "r") as zip_ref:
|
||||
zip_ref.extractall(args.local_dir)
|
||||
|
||||
print(f"✅ Done: {args.repo_id} -> {args.local_dir}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@ -1,43 +0,0 @@
|
||||
"""Download and extract FlashRAG index."""
|
||||
|
||||
import os
|
||||
import zipfile
|
||||
|
||||
import requests
|
||||
from tqdm import tqdm
|
||||
|
||||
from config import DATA_DIR
|
||||
|
||||
# Constants
|
||||
URL = "https://www.modelscope.cn/datasets/hhjinjiajie/FlashRAG_Dataset/resolve/master/retrieval_corpus/wiki18_100w_e5_index.zip"
|
||||
ZIP_NAME = "wiki18_100w_e5_index.zip"
|
||||
zip_path = DATA_DIR / ZIP_NAME
|
||||
|
||||
# Download with progress bar
|
||||
print("📥 Downloading index...")
|
||||
response = requests.get(URL, stream=True)
|
||||
total_size = int(response.headers.get("content-length", 0))
|
||||
|
||||
with (
|
||||
open(zip_path, "wb") as f,
|
||||
tqdm(
|
||||
desc=ZIP_NAME,
|
||||
total=total_size,
|
||||
unit="iB",
|
||||
unit_scale=True,
|
||||
unit_divisor=1024,
|
||||
) as bar,
|
||||
):
|
||||
for data in response.iter_content(chunk_size=1024):
|
||||
size = f.write(data)
|
||||
bar.update(size)
|
||||
|
||||
# Extract
|
||||
print("📦 Extracting index...")
|
||||
with zipfile.ZipFile(zip_path, "r") as zip_ref:
|
||||
zip_ref.extractall(DATA_DIR)
|
||||
|
||||
# Clean up zip
|
||||
os.remove(zip_path)
|
||||
print("✅ Download and extraction completed successfully!")
|
||||
print(f"Index file is at: {DATA_DIR}/data00/jiajie_jin/flashrag_indexes/wiki_dpr_100w/e5_flat_inner.index")
|
||||
@ -1,54 +0,0 @@
|
||||
import argparse
|
||||
import os
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
from config import GENERATOR_MODEL_DIR, GENERATOR_MODEL_REPO_ID
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
"""Parse command line arguments.
|
||||
|
||||
Returns:
|
||||
argparse.Namespace: Parsed arguments
|
||||
"""
|
||||
parser = argparse.ArgumentParser(description="Download model from HuggingFace Hub")
|
||||
parser.add_argument(
|
||||
"--repo-id",
|
||||
type=str,
|
||||
default=GENERATOR_MODEL_REPO_ID,
|
||||
help="HuggingFace repository ID",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--local-dir",
|
||||
type=str,
|
||||
default=GENERATOR_MODEL_DIR,
|
||||
help="Local directory to save model",
|
||||
)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
"""Main function to download model."""
|
||||
args = parse_args()
|
||||
load_dotenv(override=True)
|
||||
|
||||
# Configuration
|
||||
HF_TOKEN = os.getenv("HF_TOKEN")
|
||||
|
||||
print("Downloading model to", args.local_dir)
|
||||
|
||||
# Download the model
|
||||
snapshot_download(
|
||||
token=HF_TOKEN,
|
||||
repo_id=args.repo_id,
|
||||
local_dir=args.local_dir,
|
||||
repo_type="model",
|
||||
)
|
||||
print(f"✅ Done: {args.repo_id} -> {args.local_dir}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@ -1,54 +0,0 @@
|
||||
import argparse
|
||||
import os
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
from config import RETRIEVER_MODEL_DIR, RETRIEVER_MODEL_REPO_ID
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
"""Parse command line arguments.
|
||||
|
||||
Returns:
|
||||
argparse.Namespace: Parsed arguments
|
||||
"""
|
||||
parser = argparse.ArgumentParser(description="Download model from HuggingFace Hub")
|
||||
parser.add_argument(
|
||||
"--repo-id",
|
||||
type=str,
|
||||
default=RETRIEVER_MODEL_REPO_ID,
|
||||
help="HuggingFace repository ID",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--local-dir",
|
||||
type=str,
|
||||
default=RETRIEVER_MODEL_DIR,
|
||||
help="Local directory to save model",
|
||||
)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
"""Main function to download model."""
|
||||
args = parse_args()
|
||||
load_dotenv(override=True)
|
||||
|
||||
# Configuration
|
||||
HF_TOKEN = os.getenv("HF_TOKEN")
|
||||
|
||||
print("Downloading model to", args.local_dir)
|
||||
|
||||
# Download the model
|
||||
snapshot_download(
|
||||
token=HF_TOKEN,
|
||||
repo_id=args.repo_id,
|
||||
local_dir=args.local_dir,
|
||||
repo_type="model",
|
||||
)
|
||||
print(f"✅ Done: {args.repo_id} -> {args.local_dir}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@ -1,9 +0,0 @@
|
||||
# ------------------------------------------------Environment Settings------------------------------------------------#
|
||||
gpu_id: "0"
|
||||
|
||||
# -------------------------------------------------Retrieval Settings------------------------------------------------#
|
||||
# If set the name, the model path will be find in global paths
|
||||
retrieval_method: "e5" # name or path of the retrieval model.
|
||||
index_path: "/mnt/nas/thinhlpg/data/data00/jiajie_jin/flashrag_indexes/wiki_dpr_100w/e5_flat_inner.index" # path to the indexed file
|
||||
faiss_gpu: False # whether use gpu to hold index
|
||||
corpus_path: "/mnt/nas/thinhlpg/code/DeepSearch/data/flashrag_datasets/wiki18_100w.jsonl" # path to corpus in '.jsonl' format that store the documents
|
||||
@ -1,127 +0,0 @@
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
# Add project root to sys.path to allow importing config
|
||||
# Assuming the script is at DeepSearch/scripts/serving/serve_generator.py
|
||||
# The project root (DeepSearch) is parents[2]
|
||||
PROJ_ROOT = Path(__file__).resolve().parents[2]
|
||||
if str(PROJ_ROOT) not in sys.path:
|
||||
sys.path.append(str(PROJ_ROOT))
|
||||
|
||||
# Import after adjusting sys.path
|
||||
try:
|
||||
from config import (
|
||||
GENERATOR_MODEL_REPO_ID,
|
||||
GENERATOR_SERVER_PORT,
|
||||
MODEL_CONFIG,
|
||||
logger,
|
||||
)
|
||||
except ImportError as e:
|
||||
# Use print here as logger might not be available if import failed
|
||||
print(
|
||||
f"Error importing config: {e}. Make sure config.py is in the project root ({PROJ_ROOT}) and added to sys.path."
|
||||
)
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
def launch_sglang_server(
|
||||
model_id: str,
|
||||
port: int,
|
||||
context_length: int,
|
||||
host: str = "0.0.0.0",
|
||||
dtype: str = "bfloat16",
|
||||
) -> None:
|
||||
"""Launches the SGLang server using specified configurations.
|
||||
|
||||
Args:
|
||||
model_id: The Hugging Face repository ID of the model.
|
||||
port: The port number for the server.
|
||||
context_length: The maximum context length for the model.
|
||||
host: The host address for the server.
|
||||
dtype: The data type for the model (e.g., 'bfloat16', 'float16').
|
||||
"""
|
||||
command = [
|
||||
sys.executable, # Use the current Python interpreter
|
||||
"-m",
|
||||
"sglang.launch_server",
|
||||
"--model-path",
|
||||
model_id,
|
||||
"--context-length",
|
||||
str(context_length),
|
||||
"--enable-metrics",
|
||||
"--dtype",
|
||||
dtype,
|
||||
"--host",
|
||||
host,
|
||||
"--port",
|
||||
str(port),
|
||||
"--mem-fraction-static",
|
||||
"0.5",
|
||||
"--trust-remote-code",
|
||||
# Recommended by SGLang for stability sometimes
|
||||
"--disable-overlap",
|
||||
# Can sometimes cause issues
|
||||
"--disable-radix-cache",
|
||||
]
|
||||
|
||||
# Log the command clearly
|
||||
command_str = " ".join(command)
|
||||
logger.info(f"🚀 Launching SGLang server with command: {command_str}")
|
||||
|
||||
process = None # Initialize process to None
|
||||
try:
|
||||
# Use Popen to start the server process
|
||||
# It runs in the foreground relative to this script,
|
||||
# but allows us to catch KeyboardInterrupt cleanly.
|
||||
process = subprocess.Popen(command)
|
||||
# Wait for the process to complete (e.g., user interruption)
|
||||
process.wait()
|
||||
# Check return code after waiting
|
||||
if process.returncode != 0:
|
||||
logger.error(f"💥 SGLang server process exited with error code: {process.returncode}")
|
||||
sys.exit(process.returncode)
|
||||
else:
|
||||
logger.info("✅ SGLang server process finished gracefully.")
|
||||
|
||||
except FileNotFoundError:
|
||||
logger.error(f"💥 Error: Python executable or sglang module not found.")
|
||||
logger.error(f"Ensure '{sys.executable}' is correct and sglang is installed.")
|
||||
sys.exit(1)
|
||||
except KeyboardInterrupt:
|
||||
logger.info("🛑 SGLang server launch interrupted by user. Stopping server...")
|
||||
# Attempt to terminate the process gracefully
|
||||
if process and process.poll() is None: # Check if process exists and is running
|
||||
process.terminate()
|
||||
try:
|
||||
process.wait(timeout=5) # Wait a bit for termination
|
||||
logger.info("✅ Server terminated gracefully.")
|
||||
except subprocess.TimeoutExpired:
|
||||
logger.warning("⚠️ Server did not terminate gracefully, forcing kill.")
|
||||
process.kill()
|
||||
sys.exit(0) # Exit cleanly after interrupt
|
||||
except Exception as e:
|
||||
# Catch any other unexpected exceptions during launch or waiting
|
||||
logger.error(f"🚨 An unexpected error occurred: {e}")
|
||||
# Ensure process is cleaned up if it exists
|
||||
if process and process.poll() is None:
|
||||
process.kill()
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Get context length from config, default to 8192
|
||||
context_len = MODEL_CONFIG.get("max_seq_length", 8192)
|
||||
|
||||
logger.info("----------------------------------------------------")
|
||||
logger.info("✨ Starting SGLang Generator Server ✨")
|
||||
logger.info(f" Model ID: {GENERATOR_MODEL_REPO_ID}")
|
||||
logger.info(f" Port: {GENERATOR_SERVER_PORT}")
|
||||
logger.info(f" Context Length: {context_len}")
|
||||
logger.info("----------------------------------------------------")
|
||||
|
||||
launch_sglang_server(
|
||||
model_id=GENERATOR_MODEL_REPO_ID,
|
||||
port=GENERATOR_SERVER_PORT,
|
||||
context_length=context_len,
|
||||
)
|
||||
@ -1,113 +0,0 @@
|
||||
import argparse
|
||||
import asyncio
|
||||
from collections import deque
|
||||
from typing import List, Tuple, Union
|
||||
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from flashrag.config import Config
|
||||
from flashrag.utils import get_retriever
|
||||
from pydantic import BaseModel
|
||||
|
||||
from config import RETRIEVER_SERVER_PORT
|
||||
|
||||
app = FastAPI()
|
||||
|
||||
retriever_list = []
|
||||
available_retrievers = deque()
|
||||
retriever_semaphore = None
|
||||
|
||||
|
||||
def init_retriever(args):
|
||||
global retriever_semaphore
|
||||
config = Config(args.config)
|
||||
for i in range(args.num_retriever):
|
||||
print(f"Initializing retriever {i + 1}/{args.num_retriever}")
|
||||
retriever = get_retriever(config)
|
||||
retriever_list.append(retriever)
|
||||
available_retrievers.append(i)
|
||||
# create a semaphore to limit the number of retrievers that can be used concurrently
|
||||
retriever_semaphore = asyncio.Semaphore(args.num_retriever)
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
async def health_check():
|
||||
return {"status": "healthy", "retrievers": {"total": len(retriever_list), "available": len(available_retrievers)}}
|
||||
|
||||
|
||||
class QueryRequest(BaseModel):
|
||||
query: str
|
||||
top_n: int = 10
|
||||
return_score: bool = False
|
||||
|
||||
|
||||
class BatchQueryRequest(BaseModel):
|
||||
query: List[str]
|
||||
top_n: int = 10
|
||||
return_score: bool = False
|
||||
|
||||
|
||||
class Document(BaseModel):
|
||||
id: str
|
||||
contents: str
|
||||
|
||||
|
||||
@app.post("/search", response_model=Union[Tuple[List[Document], List[float]], List[Document]])
|
||||
async def search(request: QueryRequest):
|
||||
query = request.query
|
||||
top_n = request.top_n
|
||||
return_score = request.return_score
|
||||
|
||||
if not query or not query.strip():
|
||||
print(f"Query content cannot be empty: {query}")
|
||||
raise HTTPException(status_code=400, detail="Query content cannot be empty")
|
||||
|
||||
async with retriever_semaphore:
|
||||
retriever_idx = available_retrievers.popleft()
|
||||
try:
|
||||
if return_score:
|
||||
results, scores = retriever_list[retriever_idx].search(query, top_n, return_score)
|
||||
return [Document(id=result["id"], contents=result["contents"]) for result in results], scores
|
||||
else:
|
||||
results = retriever_list[retriever_idx].search(query, top_n, return_score)
|
||||
return [Document(id=result["id"], contents=result["contents"]) for result in results]
|
||||
finally:
|
||||
available_retrievers.append(retriever_idx)
|
||||
|
||||
|
||||
@app.post("/batch_search", response_model=Union[List[List[Document]], Tuple[List[List[Document]], List[List[float]]]])
|
||||
async def batch_search(request: BatchQueryRequest):
|
||||
query = request.query
|
||||
top_n = request.top_n
|
||||
return_score = request.return_score
|
||||
|
||||
async with retriever_semaphore:
|
||||
retriever_idx = available_retrievers.popleft()
|
||||
try:
|
||||
if return_score:
|
||||
results, scores = retriever_list[retriever_idx].batch_search(query, top_n, return_score)
|
||||
return [
|
||||
[Document(id=result["id"], contents=result["contents"]) for result in results[i]]
|
||||
for i in range(len(results))
|
||||
], scores
|
||||
else:
|
||||
results = retriever_list[retriever_idx].batch_search(query, top_n, return_score)
|
||||
return [
|
||||
[Document(id=result["id"], contents=result["contents"]) for result in results[i]]
|
||||
for i in range(len(results))
|
||||
]
|
||||
finally:
|
||||
available_retrievers.append(retriever_idx)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--config", type=str, default="./retriever_config.yaml")
|
||||
parser.add_argument("--num_retriever", type=int, default=1)
|
||||
parser.add_argument("--port", type=int, default=RETRIEVER_SERVER_PORT)
|
||||
args = parser.parse_args()
|
||||
|
||||
init_retriever(args)
|
||||
|
||||
import uvicorn
|
||||
|
||||
uvicorn.run(app, host="0.0.0.0", port=args.port)
|
||||
@ -1,135 +0,0 @@
|
||||
import json
|
||||
import math # Import math for ceiling division
|
||||
import sys
|
||||
import traceback # Import traceback
|
||||
from pathlib import Path
|
||||
|
||||
import pandas as pd
|
||||
|
||||
# Add project root to Python path if needed (adjust relative path as necessary)
|
||||
project_root = Path(__file__).resolve().parent.parent
|
||||
sys.path.append(str(project_root))
|
||||
|
||||
from src.embeddings import CustomHuggingFaceEmbeddings
|
||||
|
||||
# Import FAISS after potentially adding to sys.path
|
||||
try:
|
||||
from langchain_community.vectorstores import FAISS
|
||||
except ImportError:
|
||||
print("Error: langchain_community or FAISS not installed. Please install with 'pip install langchain faiss-cpu'")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
def build_faiss_index_from_csv(csv_path: str, index_save_path: str, batch_size: int = 128) -> None:
|
||||
"""Builds a FAISS index from a CSV containing paragraph content and metadata.
|
||||
|
||||
Reads a CSV file, generates embeddings for the 'content' column in batches,
|
||||
and saves the FAISS index files (index.faiss, index.pkl) locally.
|
||||
|
||||
Args:
|
||||
csv_path: Path to the input CSV file (e.g., data/processed/paragraphs.csv).
|
||||
index_save_path: Path to the directory where the index files should be saved.
|
||||
batch_size: Number of texts to process in each embedding batch.
|
||||
"""
|
||||
print(f"Loading paragraphs from {csv_path}")
|
||||
try:
|
||||
df = pd.read_csv(csv_path)
|
||||
except FileNotFoundError:
|
||||
print(f"Error: CSV file not found at {csv_path}. Please run the extraction script first.")
|
||||
return
|
||||
except Exception as e:
|
||||
print(f"Error reading CSV file: {e}")
|
||||
return
|
||||
|
||||
if "content" not in df.columns or "metadata" not in df.columns:
|
||||
print("Error: CSV file must contain 'content' and 'metadata' columns.")
|
||||
return
|
||||
|
||||
if df.empty:
|
||||
print("Warning: Input CSV file is empty. No index will be built.")
|
||||
return
|
||||
|
||||
# Prepare documents for FAISS
|
||||
texts = df["content"].astype(str).tolist()
|
||||
metadatas = []
|
||||
try:
|
||||
metadatas = [json.loads(m) for m in df["metadata"].tolist()]
|
||||
print(f"Prepared {len(texts)} texts and {len(metadatas)} metadatas.")
|
||||
except json.JSONDecodeError as e:
|
||||
print(f"Error parsing metadata JSON: {e}. Check the format in {csv_path}")
|
||||
traceback.print_exc() # Print traceback for JSON errors
|
||||
return
|
||||
except Exception as e:
|
||||
print(f"Error processing metadata: {e}")
|
||||
traceback.print_exc() # Print traceback for other metadata errors
|
||||
return
|
||||
|
||||
if not texts or not metadatas or len(texts) != len(metadatas):
|
||||
print(f"Error: Mismatch or empty texts/metadatas. Texts: {len(texts)}, Metadatas: {len(metadatas)}")
|
||||
return
|
||||
|
||||
print("Initializing embeddings model...")
|
||||
try:
|
||||
embeddings = CustomHuggingFaceEmbeddings()
|
||||
except Exception as e:
|
||||
print(f"Error initializing embeddings model: {e}")
|
||||
traceback.print_exc()
|
||||
return
|
||||
print("Embeddings model initialized successfully.")
|
||||
|
||||
vectorstore = None
|
||||
num_batches = math.ceil(len(texts) / batch_size)
|
||||
print(f"Processing {len(texts)} texts in {num_batches} batches of size {batch_size}...")
|
||||
|
||||
for i in range(num_batches):
|
||||
start_idx = i * batch_size
|
||||
end_idx = min((i + 1) * batch_size, len(texts))
|
||||
batch_texts = texts[start_idx:end_idx]
|
||||
batch_metadatas = metadatas[start_idx:end_idx]
|
||||
print(f" Processing batch {i + 1}/{num_batches} (indices {start_idx}-{end_idx - 1})...")
|
||||
|
||||
try:
|
||||
if i == 0:
|
||||
# Initialize the vector store with the first batch
|
||||
print(f" Initializing FAISS index with first batch...")
|
||||
vectorstore = FAISS.from_texts(texts=batch_texts, embedding=embeddings, metadatas=batch_metadatas)
|
||||
print(" FAISS index initialized.")
|
||||
else:
|
||||
# Add subsequent batches to the existing store
|
||||
if vectorstore is None:
|
||||
print("Error: vectorstore is None after first batch, cannot add more texts.")
|
||||
return # Should not happen if first batch succeeded
|
||||
print(f" Adding batch {i + 1} to FAISS index...")
|
||||
vectorstore.add_texts(texts=batch_texts, metadatas=batch_metadatas)
|
||||
print(f" Batch {i + 1} added.")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error processing batch {i + 1} (indices {start_idx}-{end_idx - 1}): {e}")
|
||||
traceback.print_exc()
|
||||
print("Stopping index creation due to error in batch processing.")
|
||||
return # Exit if any batch fails
|
||||
|
||||
if vectorstore is None:
|
||||
print("Error: Failed to create or add any data to the vectorstore.")
|
||||
return
|
||||
|
||||
# Save the completed index
|
||||
try:
|
||||
print(f"Attempting to save final FAISS index files to directory: {index_save_path}")
|
||||
# Ensure the target directory exists before saving
|
||||
Path(index_save_path).mkdir(parents=True, exist_ok=True)
|
||||
vectorstore.save_local(index_save_path)
|
||||
print(f"Successfully saved final FAISS index files (index.faiss, index.pkl) to: {index_save_path}")
|
||||
except Exception as e:
|
||||
print(f"Error during final vectorstore.save_local to {index_save_path}: {e}")
|
||||
traceback.print_exc()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Define paths relative to this script or use absolute paths
|
||||
PROCESSED_DIR = Path("data/processed")
|
||||
INPUT_CSV = str(PROCESSED_DIR / "paragraphs.csv")
|
||||
# FAISS save_local will save index.faiss and index.pkl in this directory
|
||||
INDEX_SAVE_DIR = str(PROCESSED_DIR) # Save directly to processed dir
|
||||
|
||||
build_faiss_index_from_csv(INPUT_CSV, INDEX_SAVE_DIR, batch_size=128)
|
||||
@ -1,30 +0,0 @@
|
||||
#!/bin/bash
|
||||
# This script is taken from https://github.com/StonyBrookNLP/musique with slight modifications
|
||||
|
||||
set -e
|
||||
set -x
|
||||
|
||||
# If gdown doesn't work, you can download files from mentioned URLs manually
|
||||
# and put them at appropriate locations.
|
||||
pip install gdown
|
||||
|
||||
ZIP_NAME="musique_v1.0.zip"
|
||||
|
||||
# URL: https://drive.google.com/file/d/1tGdADlNjWFaHLeZZGShh2IRcpO6Lv24h/view?usp=sharing
|
||||
gdown --id 1tGdADlNjWFaHLeZZGShh2IRcpO6Lv24h --output $ZIP_NAME
|
||||
|
||||
TARGET_DIR="./data/raw"
|
||||
mkdir -p $TARGET_DIR
|
||||
unzip -o $(basename $ZIP_NAME) -d $TARGET_DIR # Extract directly into target
|
||||
|
||||
# Move contents from the extracted 'data' folder up one level
|
||||
mv $TARGET_DIR/data/* $TARGET_DIR/
|
||||
|
||||
# Clean up the empty directory and the zip
|
||||
rm -rf $TARGET_DIR/data
|
||||
rm $ZIP_NAME
|
||||
|
||||
# TODO: prevent these from zipping in.
|
||||
rm -rf __MACOSX
|
||||
# Clean up potential extracted .DS_Store
|
||||
rm -f $TARGET_DIR/.DS_Store
|
||||
@ -1,101 +0,0 @@
|
||||
import json
|
||||
import sys
|
||||
from collections import defaultdict # Use defaultdict for cleaner accumulation
|
||||
from pathlib import Path
|
||||
|
||||
import pandas as pd
|
||||
|
||||
# Add project root to Python path if needed (adjust relative path as necessary)
|
||||
# project_root = Path(__file__).resolve().parent.parent
|
||||
# sys.path.append(str(project_root))
|
||||
# from config import logger # Assuming you have a logger setup
|
||||
|
||||
|
||||
def extract_unique_paragraphs(input_paths: list[str], output_csv_path: str) -> None:
|
||||
"""Extracts unique paragraphs from specified JSONL files.
|
||||
|
||||
Reads Musique JSONL files (train, dev, test), finds unique paragraphs
|
||||
(regardless of is_supporting flag), combines title and text,
|
||||
tracks source question IDs, and saves to CSV.
|
||||
|
||||
Args:
|
||||
input_paths: A list of paths to the input JSONL files.
|
||||
output_csv_path: Path to save the output CSV file.
|
||||
"""
|
||||
output_dir = Path(output_csv_path).parent
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Use paragraph content as key, value is the set of source question IDs
|
||||
paragraphs_data = defaultdict(set)
|
||||
print("Starting paragraph extraction (including non-supporting)...")
|
||||
|
||||
for file_path in input_paths:
|
||||
print(f"Processing file: {file_path}")
|
||||
try:
|
||||
with open(file_path, "r", encoding="utf-8") as infile:
|
||||
for line_num, line in enumerate(infile, 1):
|
||||
try:
|
||||
data = json.loads(line)
|
||||
main_question_id = data.get("id")
|
||||
if not main_question_id:
|
||||
print(f"Warning: Missing 'id' in line {line_num} of {file_path}")
|
||||
continue
|
||||
|
||||
for p in data.get("paragraphs", []):
|
||||
title = p.get("title", "No Title")
|
||||
text = p.get("paragraph_text", "")
|
||||
content = f"{title}\n{text}".strip()
|
||||
|
||||
if not content:
|
||||
continue # Skip empty paragraphs
|
||||
|
||||
paragraphs_data[content].add(main_question_id)
|
||||
|
||||
except json.JSONDecodeError:
|
||||
print(f"Warning: Skipping invalid JSON in line {line_num} of {file_path}")
|
||||
except Exception as e:
|
||||
print(f"Warning: Error processing line {line_num} in {file_path}: {e}")
|
||||
except FileNotFoundError:
|
||||
print(f"Error: Input file not found: {file_path}")
|
||||
except Exception as e:
|
||||
print(f"Error reading file {file_path}: {e}")
|
||||
|
||||
print(f"Found {len(paragraphs_data)} unique paragraphs (supporting and non-supporting).")
|
||||
|
||||
# Prepare data for DataFrame
|
||||
output_list = []
|
||||
sorted_content = sorted(paragraphs_data.keys())
|
||||
for chunk_id, content in enumerate(sorted_content, 1):
|
||||
question_ids = paragraphs_data[content]
|
||||
metadata = {"source_question_ids": sorted(list(question_ids))}
|
||||
output_list.append(
|
||||
{
|
||||
"chunk_id": chunk_id,
|
||||
"content": content,
|
||||
"metadata": json.dumps(metadata), # Store metadata as JSON string
|
||||
}
|
||||
)
|
||||
|
||||
if not output_list:
|
||||
print("No paragraphs found to save.")
|
||||
return
|
||||
df = pd.DataFrame(output_list)
|
||||
try:
|
||||
df.to_csv(output_csv_path, index=False)
|
||||
print(f"Successfully saved unique paragraphs to {output_csv_path}")
|
||||
except Exception as e:
|
||||
print(f"Error saving CSV file: {e}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
RAW_DIR = Path("data/raw")
|
||||
PROCESSED_DIR = Path("data/processed")
|
||||
|
||||
input_files = [
|
||||
str(RAW_DIR / "musique_ans_v1.0_train.jsonl"),
|
||||
str(RAW_DIR / "musique_ans_v1.0_dev.jsonl"),
|
||||
str(RAW_DIR / "musique_ans_v1.0_test.jsonl"),
|
||||
]
|
||||
output_csv = str(PROCESSED_DIR / "paragraphs.csv")
|
||||
|
||||
extract_unique_paragraphs(input_files, output_csv)
|
||||
@ -1,155 +0,0 @@
|
||||
"""Prepares a deterministic sampled dev set (questions_dev.jsonl) from raw Musique dev data."""
|
||||
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import re
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def transform_musique_dev_data(input_path: str, output_path: str, sample_config: dict) -> None:
|
||||
"""Transforms Musique dev data with deterministic stratified sampling using uniform selection from sorted lists.
|
||||
|
||||
Reads dev data, categorizes by hop type (2, 3, 4), sorts categories by ID,
|
||||
selects N samples uniformly spaced from each sorted category based on sample_config,
|
||||
combines, sorts final list by ID, combines answers/aliases, extracts supporting paras,
|
||||
and writes the transformed data to output_path.
|
||||
|
||||
Args:
|
||||
input_path: Path to the input JSONL file (e.g., data/raw/musique_ans_v1.0_dev.jsonl).
|
||||
output_path: Path to the output JSONL file (e.g., data/processed/questions_dev.jsonl).
|
||||
sample_config: Dictionary specifying samples per hop type (e.g., {"2hop": 20, "3hop": 15, "4hop": 15}).
|
||||
"""
|
||||
output_dir = Path(output_path).parent
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
print(f"Reading all data from {input_path} for dev sampling...")
|
||||
all_data = []
|
||||
try:
|
||||
with open(input_path, "r", encoding="utf-8") as infile:
|
||||
for line_num, line in enumerate(infile, 1):
|
||||
try:
|
||||
data = json.loads(line)
|
||||
if "id" in data:
|
||||
all_data.append(data)
|
||||
else:
|
||||
print(f"Warning: Skipping line {line_num} due to missing 'id' field in {input_path}")
|
||||
except json.JSONDecodeError:
|
||||
print(f"Warning: Skipping invalid JSON in line {line_num} of {input_path}")
|
||||
except FileNotFoundError:
|
||||
print(f"Error: Input file not found at {input_path}")
|
||||
return
|
||||
except Exception as e:
|
||||
print(f"Error reading file {input_path}: {e}")
|
||||
return
|
||||
print(f"Read {len(all_data)} total samples from dev set.")
|
||||
|
||||
# Categorize data by hop count (2hop, 3hop, 4hop)
|
||||
categorized_data = defaultdict(list)
|
||||
print("Categorizing data by hop type (2, 3, 4)...")
|
||||
for data in all_data:
|
||||
q_id = data["id"]
|
||||
hop_type = None
|
||||
if q_id.startswith("2hop"):
|
||||
hop_type = "2hop"
|
||||
elif q_id.startswith("3hop"):
|
||||
hop_type = "3hop"
|
||||
elif q_id.startswith("4hop"):
|
||||
hop_type = "4hop"
|
||||
|
||||
if hop_type:
|
||||
categorized_data[hop_type].append(data)
|
||||
|
||||
# Deterministic sampling using sorting and uniform index selection
|
||||
final_sample_list = []
|
||||
total_target = sum(sample_config.values())
|
||||
print(f"Sampling deterministically via uniform selection from sorted lists to get {total_target} dev samples...")
|
||||
|
||||
for hop_type, target_count in sample_config.items():
|
||||
available_samples = categorized_data.get(hop_type, [])
|
||||
current_count = len(available_samples)
|
||||
print(f" {hop_type}: Found {current_count} samples, need {target_count}.")
|
||||
|
||||
if current_count == 0:
|
||||
continue
|
||||
|
||||
available_samples.sort(key=lambda x: x["id"])
|
||||
selected_samples_for_hop = []
|
||||
if current_count < target_count:
|
||||
print(f" Warning: Not enough samples for {hop_type}. Taking all {current_count} sorted samples.")
|
||||
selected_samples_for_hop = available_samples
|
||||
elif target_count > 0: # Ensure target_count is positive before selecting
|
||||
print(f" Selecting {target_count} samples uniformly from {current_count}...")
|
||||
# Calculate indices using integer interpretation of evenly spaced points
|
||||
indices_to_take = [
|
||||
int(i * (current_count - 1) / (target_count - 1)) if target_count > 1 else 0
|
||||
for i in range(target_count)
|
||||
] # Adjust index calc for edges
|
||||
indices_to_take = sorted(list(set(indices_to_take))) # Ensure unique indices
|
||||
# Simple fallback if uniqueness reduced count below target
|
||||
while len(indices_to_take) < target_count and len(indices_to_take) < current_count:
|
||||
next_val = indices_to_take[-1] + 1
|
||||
if next_val < current_count:
|
||||
indices_to_take.append(next_val)
|
||||
else: # Cannot add more unique indices
|
||||
break
|
||||
selected_samples_for_hop = [
|
||||
available_samples[idx] for idx in indices_to_take[:target_count]
|
||||
] # Select based on unique indices, capped at target
|
||||
|
||||
final_sample_list.extend(selected_samples_for_hop)
|
||||
|
||||
print(f"Selected {len(final_sample_list)} dev samples in total.")
|
||||
|
||||
# Sort the final combined list by ID for consistent output order
|
||||
print("Sorting the final combined dev sample list by ID...")
|
||||
final_sample_list.sort(key=lambda x: x["id"])
|
||||
|
||||
# Process and write the selected samples
|
||||
print(f"Processing and writing {len(final_sample_list)} selected dev samples to {output_path}...")
|
||||
count = 0
|
||||
try:
|
||||
with open(output_path, "w", encoding="utf-8") as outfile:
|
||||
for data in final_sample_list:
|
||||
try:
|
||||
supporting_paragraphs = [
|
||||
p["paragraph_text"] for p in data.get("paragraphs", []) if p.get("is_supporting", False)
|
||||
]
|
||||
main_answer = data.get("answer", "")
|
||||
aliases = data.get("answer_aliases", [])
|
||||
all_answers = [main_answer] + (aliases if isinstance(aliases, list) else [])
|
||||
valid_answers = [str(ans).strip() for ans in all_answers if ans and str(ans).strip()]
|
||||
unique_valid_answers = list(set(valid_answers)) # Keep unique, don't sort alphabetically
|
||||
combined_answer_str = " OR ".join(unique_valid_answers)
|
||||
|
||||
output_data = {
|
||||
"id": data.get("id"),
|
||||
"question": data.get("question"),
|
||||
"answer": combined_answer_str,
|
||||
"supporting_paragraphs": supporting_paragraphs,
|
||||
}
|
||||
outfile.write(json.dumps(output_data) + "\n")
|
||||
count += 1
|
||||
except KeyError as e:
|
||||
print(f"Skipping sample due to missing key {e}: {data.get('id')}")
|
||||
print(f"Successfully processed and wrote {count} dev samples.")
|
||||
except Exception as e:
|
||||
print(f"An unexpected error occurred during writing: {e}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Define file paths relative to the project root
|
||||
# Ensure this script is run from the project root or adjust paths accordingly
|
||||
RAW_DIR = Path("data/raw")
|
||||
PROCESSED_DIR = Path("data/processed")
|
||||
|
||||
# Define sampling configuration for the dev set
|
||||
DEV_SAMPLING_CONFIG = {"2hop": 20, "3hop": 15, "4hop": 15} # Total = 50
|
||||
|
||||
INPUT_FILE = RAW_DIR / "musique_ans_v1.0_dev.jsonl"
|
||||
OUTPUT_FILE = PROCESSED_DIR / "questions_dev.jsonl"
|
||||
|
||||
transform_musique_dev_data(str(INPUT_FILE), str(OUTPUT_FILE), DEV_SAMPLING_CONFIG)
|
||||
|
||||
print(f"\nMusique DEV JSONL transformation and deterministic sampling complete.")
|
||||
@ -1,172 +0,0 @@
|
||||
import json
|
||||
import math # Keep math import
|
||||
import os
|
||||
import re # Import re for parsing ID
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
|
||||
# import random # No longer needed
|
||||
# SEED = 42 # No longer needed
|
||||
# random.seed(SEED) # No longer needed
|
||||
|
||||
|
||||
def transform_musique_data(input_path: str, output_path: str, sample_config: dict) -> None:
|
||||
"""Transforms Musique data with deterministic stratified sampling using uniform selection from sorted lists.
|
||||
|
||||
Reads data, categorizes by detailed hop type, sorts categories by ID,
|
||||
selects N samples uniformly spaced from each sorted category,
|
||||
combines, sorts final list by ID, and writes to output.
|
||||
|
||||
Args:
|
||||
input_path: Path to the input JSONL file.
|
||||
output_path: Path to the output JSONL file.
|
||||
sample_config: Dictionary specifying samples per detailed hop type (e.g., {"2hop": 400, "3hop1": 150, ...}).
|
||||
"""
|
||||
output_dir = Path(output_path).parent
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
print(f"Reading all data from {input_path} for sampling...")
|
||||
all_data = []
|
||||
try:
|
||||
with open(input_path, "r", encoding="utf-8") as infile:
|
||||
for line_num, line in enumerate(infile, 1):
|
||||
try:
|
||||
data = json.loads(line)
|
||||
if "id" in data:
|
||||
all_data.append(data)
|
||||
else:
|
||||
print(f"Warning: Skipping line {line_num} due to missing 'id' field in {input_path}")
|
||||
except json.JSONDecodeError:
|
||||
print(f"Warning: Skipping invalid JSON in line {line_num} of {input_path}")
|
||||
except FileNotFoundError:
|
||||
print(f"Error: Input file not found at {input_path}")
|
||||
return
|
||||
except Exception as e:
|
||||
print(f"Error reading file {input_path}: {e}")
|
||||
return
|
||||
print(f"Read {len(all_data)} total samples with IDs.")
|
||||
|
||||
# Detailed Categorization by hop type
|
||||
categorized_data = defaultdict(list)
|
||||
print("Categorizing data by detailed hop type (e.g., 3hop1, 4hop2)...")
|
||||
for data in all_data:
|
||||
q_id = data["id"]
|
||||
match = re.match(r"^(2hop|3hop[12]|4hop[123])__", q_id)
|
||||
if match:
|
||||
detailed_hop_type = match.group(1)
|
||||
categorized_data[detailed_hop_type].append(data)
|
||||
# else: # Optional: log if an ID doesn't match expected pattern
|
||||
# print(f"Warning: ID {q_id} does not match expected hop pattern.")
|
||||
|
||||
# Deterministic sampling using sorting and uniform index selection
|
||||
final_sample_list = []
|
||||
total_target = sum(sample_config.values())
|
||||
print(f"Sampling deterministically via uniform selection from sorted lists to get {total_target} samples...")
|
||||
# Check if all requested hop types exist in config
|
||||
for hop_type in sample_config.keys():
|
||||
if hop_type not in categorized_data:
|
||||
print(f"Warning: Hop type '{hop_type}' requested in config but not found in data.")
|
||||
|
||||
for hop_type, target_count in sample_config.items():
|
||||
available_samples = categorized_data.get(hop_type, [])
|
||||
current_count = len(available_samples)
|
||||
print(f" {hop_type}: Found {current_count} samples, need {target_count}.")
|
||||
|
||||
if current_count == 0:
|
||||
continue
|
||||
|
||||
# Sort the list for this category by ID
|
||||
available_samples.sort(key=lambda x: x["id"])
|
||||
|
||||
selected_samples_for_hop = []
|
||||
if current_count < target_count:
|
||||
print(f" Warning: Not enough samples for {hop_type}. Taking all {current_count} sorted samples.")
|
||||
selected_samples_for_hop = available_samples
|
||||
else:
|
||||
# Select target_count indices spread uniformly across the available samples
|
||||
print(f" Selecting {target_count} samples uniformly from {current_count}...")
|
||||
# Calculate indices using integer interpretation of evenly spaced points
|
||||
indices_to_take = [int(i * current_count / target_count) for i in range(target_count)]
|
||||
# Ensure uniqueness in case of rounding issues with small numbers (though unlikely here)
|
||||
indices_to_take = sorted(list(set(indices_to_take)))
|
||||
# Adjust if rounding resulted in fewer than target_count unique indices
|
||||
while len(indices_to_take) < target_count:
|
||||
# This is a fallback, shouldn't happen if current_count >= target_count
|
||||
# Add indices from the end if needed, avoiding duplicates
|
||||
next_idx = indices_to_take[-1] + 1
|
||||
if next_idx < current_count and next_idx not in indices_to_take:
|
||||
indices_to_take.append(next_idx)
|
||||
else: # Should not be reachable if logic is sound
|
||||
break
|
||||
|
||||
# Select samples at the calculated indices
|
||||
selected_samples_for_hop = [
|
||||
available_samples[idx] for idx in indices_to_take[:target_count]
|
||||
] # Ensure we take exactly target_count
|
||||
|
||||
final_sample_list.extend(selected_samples_for_hop)
|
||||
|
||||
print(f"Selected {len(final_sample_list)} samples in total.")
|
||||
|
||||
# Sort the final combined list by ID for consistent output order
|
||||
print("Sorting the final combined sample list by ID...")
|
||||
final_sample_list.sort(key=lambda x: x["id"])
|
||||
|
||||
# Process and write the selected samples
|
||||
print(f"Processing and writing {len(final_sample_list)} selected samples to {output_path}...")
|
||||
count = 0
|
||||
try:
|
||||
with open(output_path, "w", encoding="utf-8") as outfile:
|
||||
for data in final_sample_list:
|
||||
try:
|
||||
supporting_paragraphs = [
|
||||
p["paragraph_text"] for p in data.get("paragraphs", []) if p.get("is_supporting", False)
|
||||
]
|
||||
|
||||
main_answer = data.get("answer", "")
|
||||
aliases = data.get("answer_aliases", [])
|
||||
|
||||
all_answers = [main_answer] + (aliases if isinstance(aliases, list) else [])
|
||||
valid_answers = [str(ans).strip() for ans in all_answers if ans and str(ans).strip()]
|
||||
unique_valid_answers = list(set(valid_answers))
|
||||
|
||||
combined_answer_str = " OR ".join(unique_valid_answers)
|
||||
|
||||
output_data = {
|
||||
"id": data.get("id"),
|
||||
"question": data.get("question"),
|
||||
"answer": combined_answer_str,
|
||||
"supporting_paragraphs": supporting_paragraphs,
|
||||
}
|
||||
outfile.write(json.dumps(output_data) + "\n")
|
||||
count += 1
|
||||
except KeyError as e:
|
||||
print(f"Skipping sample due to missing key {e}: {data.get('id')}")
|
||||
print(f"Successfully processed and wrote {count} samples.")
|
||||
except Exception as e:
|
||||
print(f"An unexpected error occurred during writing: {e}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Define file paths
|
||||
RAW_DIR = Path("data/raw")
|
||||
PROCESSED_DIR = Path("data/processed")
|
||||
|
||||
# Define detailed sampling configuration
|
||||
SAMPLING_CONFIG = {
|
||||
"2hop": 400,
|
||||
"3hop1": 150,
|
||||
"3hop2": 150,
|
||||
"4hop1": 100,
|
||||
"4hop2": 100,
|
||||
"4hop3": 100,
|
||||
} # Total = 1000
|
||||
|
||||
transform_musique_data(
|
||||
str(RAW_DIR / "musique_ans_v1.0_train.jsonl"), str(PROCESSED_DIR / "questions.jsonl"), SAMPLING_CONFIG
|
||||
)
|
||||
|
||||
print(
|
||||
"\nMusique JSONL transformation and detailed deterministic sampling (uniform selection from sorted) complete."
|
||||
)
|
||||
# Note: Dev/Test files are not processed by default with this sampling logic.
|
||||
@ -1 +0,0 @@
|
||||
Subproject commit 7e60ab26825a452f8ee8eb19799d1cb6c1746326
|
||||
Loading…
Reference in new issue