thinhlpg
1 month ago
parent
bd02305efb
commit
14ef79a4f5
@ -0,0 +1,659 @@
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
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 tqdm import tqdm
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from config import logger
|
||||
from src.agent import Agent, AgenticOutputs
|
||||
from src.prompts import build_user_prompt, get_system_prompt
|
||||
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 = 2):
|
||||
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", 20)
|
||||
generator_max_output_len = getattr(self.config, "generator_max_output_len", 1024)
|
||||
|
||||
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 as e:
|
||||
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", 20)
|
||||
generator_max_output_len = getattr(config, "generator_max_output_len", 1024)
|
||||
|
||||
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)
|
||||
@ -0,0 +1,68 @@
|
||||
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()
|
||||
@ -0,0 +1,43 @@
|
||||
"""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")
|
||||
@ -0,0 +1,54 @@
|
||||
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()
|
||||
@ -0,0 +1,54 @@
|
||||
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()
|
||||
@ -0,0 +1,9 @@
|
||||
# ------------------------------------------------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
|
||||
@ -0,0 +1,125 @@
|
||||
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),
|
||||
"--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,
|
||||
)
|
||||
@ -0,0 +1,113 @@
|
||||
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)
|
||||
Loading…
Reference in new issue