diff --git a/learn2rag/compose/__init__.py b/learn2rag/compose/__init__.py index 688c969..1877c1c 100644 --- a/learn2rag/compose/__init__.py +++ b/learn2rag/compose/__init__.py @@ -206,8 +206,15 @@ def start(self) -> None: try: for file in self.content.get('files', []): file_path = Path(file['path']).expanduser().absolute() - file_path.parent.mkdir(parents=True, exist_ok=True) - file_path.write_text(file['content']) + if not file_path.exists() or file.get('force', True): + file_path.parent.mkdir(parents=True, exist_ok=True) + if 'content' in file: + content = file['content'] + elif 'src' in file: + content = Path(file['src']).read_text() + else: + raise NotImplementedError(file) + file_path.write_text(content) except Exception as e: con.rollback() raise e diff --git a/learn2rag/evaluation/example.py b/learn2rag/evaluation/example.py index 2343a65..841d59a 100644 --- a/learn2rag/evaluation/example.py +++ b/learn2rag/evaluation/example.py @@ -51,5 +51,6 @@ def process_qa(dataset_name: str, qa_rows: Any) -> None: # qa_rows = read_dataset_qa('hotpot_qa', 'distractor', 'validation') # process_qa('hotpot_qa', qa_rows.select(range(3))) - qa_rows = read_dataset_qa('WikiEval', '', 'train') + target_path = pathlib.Path('./datasets/WikiEval/source') + qa_rows = read_dataset_qa(target_path, split='train') process_qa('WikiEval', qa_rows) diff --git a/learn2rag/evaluation/tools.py b/learn2rag/evaluation/tools.py index b927bc1..7e7c8c3 100644 --- a/learn2rag/evaluation/tools.py +++ b/learn2rag/evaluation/tools.py @@ -6,7 +6,8 @@ import logging import datasets import json_stream # type: ignore[import-untyped] -from typing import Any, Callable +from typing import Any, Callable, cast +import pandas as pd import learn2rag.pipeline.ingestion from learn2rag.pipeline.config import opt_config @@ -78,13 +79,39 @@ def ingest_dataset_documents(dataset_name: str) -> None: # FIXME # learn2rag.pipeline.ingestion.index(user_config, opt_config) +def read_dataset_qa(file_path: pathlib.Path | str, split: str | None = None) -> Any: + target_path = pathlib.Path(file_path) + logging.debug(f'Loading dataset from: {target_path}') + + if not target_path.exists(): + raise FileNotFoundError(f"Dataset path does not exist: {target_path}") + + if target_path.suffix.lower() == '.csv': + logging.debug('load csv file') + with open(target_path, 'r', encoding='utf-8') as f: + first_line = f.readline() + try: + # Sniff for comma, semicolon, or tab + dialect = csv.Sniffer().sniff(first_line, delimiters=",;\t") + detected_sep = dialect.delimiter + except csv.Error: + # Fallback to standard comma if sniffing fails + detected_sep = ',' + + logging.debug(f'load csv file with delimiter: "{detected_sep}"') + dataset_dict = datasets.load_dataset('csv', data_files=str(target_path), sep=detected_sep) + else: + logging.debug(f'load HF dataset ') + dataset_dict = datasets.load_from_disk(str(target_path)) -def read_dataset_qa(dataset_name: str, subdirectory: str, split: str | None=None) -> Any: - logging.debug(f'{dataset_name=}') - dataset_work_dir = pathlib.Path('./datasets') / dataset_name - dataset_dict = datasets.load_from_disk(dataset_work_dir / 'source' / subdirectory) return dataset_dict[split] if split is not None else dataset_dict +# def read_dataset_qa(dataset_name: str, subdirectory: str, split: str | None=None) -> Any: +# logging.debug(f'{dataset_name=}') +# dataset_work_dir = pathlib.Path('./datasets') / dataset_name +# dataset_dict = datasets.load_from_disk(dataset_work_dir / 'source' / subdirectory) +# return dataset_dict[split] if split is not None else dataset_dict + def basic_pipeline(dataset_name: str, question: str) -> dict[str, Any]: user_config = { diff --git a/learn2rag/optimization/__init__.py b/learn2rag/optimization/__init__.py new file mode 100644 index 0000000..7fd91ad --- /dev/null +++ b/learn2rag/optimization/__init__.py @@ -0,0 +1,98 @@ +import logging +import pathlib +import os +import json +import collections.abc +import copy +import yaml +import argparse +from typing import Any, Dict, Callable +from collections.abc import Mapping + +from . import baseline_optimization +from . import retrieval_optimization + +OPTIMIZATION_STRATEGIES: Dict[str, Callable[..., Any]] = { + "baseline": baseline_optimization.run, + "retrieval": retrieval_optimization.run, +} + + +#TODO : now we need to copy the dataset to here manually it should consider in installation maybe ! +# {storage_path}/datasets/WikiEval/ +def main() -> None: + parser = argparse.ArgumentParser(description="Run RAG baseline optimization.") + parser.add_argument("task", help="Module task name (e.g., learn2rag.optimization)") + parser.add_argument("--logging-config", type=str, help="Path to logging config yml") + parser.add_argument("--registry-path", type=str, help="Path to registry.json") + parser.add_argument("--dataset", type=str, default="WikiEval") + parser.add_argument("--questions", type=int, default=10) + parser.add_argument("--trials", type=int, default=10) + parser.add_argument( + "--strategy", + type=str, + default="baseline", + choices=list(OPTIMIZATION_STRATEGIES.keys()), + help="Which optimization algorithm to run" + ) + args, unknown = parser.parse_known_args() + + if args.logging_config and pathlib.Path(args.logging_config).exists(): + with open(args.logging_config, 'r') as f: + config = yaml.safe_load(f) + logging.config.dictConfig(config) + + logger = logging.getLogger(__name__) + logger.info(f"Running task: {args.task}") + logger.info(f"Optimization started for {args.dataset} using strategy: {args.strategy}") + + logging.info(f"Optimization started for {args.dataset}") + #TODO : read or get dataset_name and maxquestions and n_trails + + output_dir = pathlib.Path("./optimization/output/") + output_dir.mkdir(parents=True, exist_ok=True) + + run_optimization = OPTIMIZATION_STRATEGIES[args.strategy] + run_optimization(args.dataset, args.questions, args.trials, output_dir, args.registry_path) + + logging.info("optimization is done") + results_path = output_dir /args.dataset/ "optimization_results.json" + logging.info(f"save optimized results here : {results_path}") + if not results_path.exists(): + logging.warning(f"Optimization results not found at {results_path}") + return + + with open(results_path, "r") as f: + full_results = json.load(f) + + best_config = full_results.get("best_config", {}) + + target_config_path = os.environ.get("PIPELINE_OPT_CONFIG", "learn2rag/pipeline/opt_config.json") + + existing_config = {} + if os.path.exists(target_config_path): + try: + with open(target_config_path, "r", encoding="utf-8") as f: + existing_config = json.load(f) + except json.JSONDecodeError: + logging.error(f"Existing config at {target_config_path} is corrupted. It will be overwritten.") + + updated_config = deep_update(copy.deepcopy(existing_config), best_config) + + pathlib.Path(target_config_path).parent.mkdir(parents=True, exist_ok=True) + with open(target_config_path, "w", encoding="utf-8") as f: + json.dump(updated_config, f, indent=4) + + logging.info(f"Successfully updated opt_config at: {target_config_path}") + +def deep_update(source: dict[str, Any], overrides:Mapping[str, Any]) -> dict[str, Any]: + """Recursively updates a dictionary.""" + for key, value in overrides.items(): + if isinstance(value, collections.abc.Mapping) and key in source: + deep_update(source.get(key, {}), value) + else: + source[key] = value + return source + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/learn2rag/optimization/baseline_optimization.py b/learn2rag/optimization/baseline_optimization.py new file mode 100644 index 0000000..ab6021e --- /dev/null +++ b/learn2rag/optimization/baseline_optimization.py @@ -0,0 +1,288 @@ +""" +RAG Pipeline Optimization with BERTScore evaluation. +""" + +import argparse +import json +import logging +import pathlib +import time +import copy +import os +from typing import Dict, Any, List, Union, Tuple, cast +from qdrant_client.models import ScoredPoint +import numpy as np +from bert_score import score as bert_score # type: ignore +from ConfigSpace import ConfigurationSpace, Integer, Categorical, ForbiddenGreaterThanRelation, Configuration +from smac import HyperparameterOptimizationFacade, Scenario + +from learn2rag.evaluation.tools import read_dataset_qa +from learn2rag.pipeline.config import opt_config +import learn2rag.pipeline.search +import learn2rag.pipeline.generate + +def load_registry(registry_source: Union[str, pathlib.Path, Dict[str, Any]] = "registry.json") -> dict[str, Any]: + if isinstance(registry_source, dict): + return registry_source + + + p = pathlib.Path(registry_source) + if not p.is_file(): + logging.error("registry file not found") + with p.open() as f: + return cast(Dict[str, Any], json.load(f)) + +def run_pipeline(question: str, user_config: Dict[str, Any], working_config: Dict[str, Any]) -> Tuple[str, str, float, float]: + t0 = time.time() + docs = learn2rag.pipeline.search.search(question, user_config, working_config) + search_time = time.time() - t0 + + t0 = time.time() + answer = learn2rag.pipeline.generate.generate(question, docs.points, working_config) + gen_time = time.time() - t0 + + doc_list = docs.points if hasattr(docs, "points") else docs + context = "" + if doc_list: + context_parts = [] + for d in doc_list: + payload = getattr(d, "payload", {}) or {} + path = payload.get("path", "unknown") if isinstance(payload, dict) else "unknown" + content = payload.get("content", "") if isinstance(payload, dict) else "" + context_parts.append(f"Source: {path}\nContent: {content}") + context = "\n\n".join(context_parts) + return answer, context[:3000], search_time, gen_time + +#I removed seed because there are no use for it +# removed dataset_name becuase it just use in yser config and now we inject it +def objective(config: Configuration, + questions: List[Dict[str, Any]], + dataset_name: str, + state: Dict[str, Any], + answers_dir: pathlib.Path + ,prompt_map: Dict[str, Any] +) -> float: + state["trial_count"] += 1 + tid = state["trial_count"] + cfg = dict(config) + logging.info(f"Trial {tid}: {cfg}") + + working_cfg = copy.deepcopy(opt_config) + working_cfg.update({ + "top_k": cfg["top_k"], + # "chunk_size": cfg["chunk_size"], + # "chunk_overlap": cfg["chunk_overlap"], + "prompt": prompt_map[cfg["prompt_template"]], + }) + + ucfg = { + "file_path": None, + "collection_name": dataset_name, + "imported_documents_file_path": None, + "llm": None, + } + env_user_cfg = os.environ.get("PIPELINE_USER_CONFIG") + if env_user_cfg and pathlib.Path(env_user_cfg).exists(): + ucfg.update(json.loads(pathlib.Path(env_user_cfg).read_text())) + + predictions, goldens = [], [] + qa_pairs = [] + t_start = time.time() + t_search, t_gen = 0.0, 0.0 + + for q in questions: + # Preserve the original “skip empty question” guard + if not q.get("question"): + continue + try: + answer, context, t_s, t_g = run_pipeline(q["question"], ucfg, working_cfg) + t_search += t_s + t_gen += t_g + predictions.append(answer) + goldens.append(q["ground_truth"]) + qa_pairs.append({**q, "generated_answer": answer, "retrieved_context": context}) + except Exception as e: + # Same behaviour as the old version: record a blank answer. TODO : check if we need this + logging.warning(f"Trial {tid}, q{q.get('id','?')} failed: {e}") + predictions.append("") + goldens.append(q["ground_truth"]) + qa_pairs.append({**q, "generated_answer": "", "retrieved_context": ""}) + + if not predictions: + return 1.0 + + t_score = time.time() + _, _, F1_gold = bert_score(predictions, goldens, lang="en", verbose=False, rescale_with_baseline=True) + scoring_time = time.time() - t_score + + # objective function + bert_gold = [max(0.0, f.item()) for f in F1_gold] + avg_bert_gold = np.mean(bert_gold) + cost = 1.0 - avg_bert_gold + total_time = time.time() - t_start + + trial_answers = { + "trial_id": tid, + "config": cfg, + "cost": float(cost), + "avg_bertscore_golden": float(avg_bert_gold), + "qa_pairs": qa_pairs, + } + answers_file = answers_dir / f"trial_{tid}_answers.json" + with open(answers_file, "w") as f: + json.dump(trial_answers, f, indent=2, default=str) + + state["best_cost"] = min(state["best_cost"], cost) + state["convergence"].append({"trial": tid, "cost": float(cost), "best_cost": float(state["best_cost"])}) + state["history"].append({ + "trial_id": tid, "config": cfg, + "avg_bertscore_golden": float(avg_bert_gold), + "cost": float(cost), "time_s": round(total_time, 2), + "search_s": round(t_search, 2), "gen_s": round(t_gen, 2), + "scoring_s": round(scoring_time, 2), + }) + + logging.info( + f"Trial {tid}: bertscore_golden={avg_bert_gold:.4f} cost={cost:.4f} " + f"time={total_time:.1f}s (search={t_search:.1f} gen={t_gen:.1f} score={scoring_time:.1f})" + ) + return float(cost) + + +def param_importance(smac: HyperparameterOptimizationFacade, output_path: pathlib.Path) -> Dict[str, Any]: + params = list(smac.scenario.configspace.keys()) + configs, costs = [], [] + for key, val in smac.runhistory.items(): + configs.append(dict(smac.runhistory.get_config(key.config_id))) + costs.append(val.cost) + if len(configs) < 3: + return {} + + raw = {} + for p in params: + groups : Dict[str, List[float]] = {} + for c, cost in zip(configs, np.array(costs)): + groups.setdefault(str(c[p]), []).append(cost) + means = [np.mean(g) for g in groups.values()] + raw[p] = float(np.var(means)) if len(means) > 1 else 0.0 + + total = sum(raw.values()) + imp = {p: round(v / total, 4) for p, v in raw.items()} if total > 0 else raw + ranking = sorted(imp, key=lambda k: imp[k], reverse=True) + result = {"method": "variance_based", "ranking": ranking, "individual": imp} + with open(output_path / "parameter_importance.json", "w") as f: + json.dump(result, f, indent=2) + return result + + +def run(dataset_name: str, max_questions: int, n_trials: int, output_dir: Union[str, pathlib.Path],registry_path: Union[str, pathlib.Path, Dict[str, Any]] ) -> Tuple[ + Dict[str, Any], List[Any], Dict[str, Any]]: + logging.info(f"registry_path is : {registry_path}") + registry = load_registry(registry_path) + datasets = registry["datasets"] + if dataset_name not in datasets: + raise ValueError(f"Unknown dataset: {dataset_name}. Available: {list(datasets.keys())}") + dcfg = datasets[dataset_name] + fields = dcfg["fields"] + base_path = pathlib.Path(dcfg["path"]) + + out = pathlib.Path(output_dir) / dataset_name + out.mkdir(parents=True, exist_ok=True) + answers_dir = out / "trial_answers" + answers_dir.mkdir(parents=True, exist_ok=True) + if base_path.suffix.lower() == '.csv': + target_path = base_path + else: + target_path = base_path / 'source' / dcfg.get("subdirectory", "") + + logging.debug(f"target path for dataset is {target_path} ") + + qa = read_dataset_qa(target_path, split=dcfg["split"]) + if max_questions: + qa = qa.select(range(min(max_questions, len(qa)))) + + questions = [ + { + "question": r[fields["q"]], + "ground_truth": r[fields["a"]], + "id": r.get(fields["id"], str(i)), + } + for i, r in enumerate(qa) + ] + prompt_map = registry["prompts"] + + cs = ConfigurationSpace(seed=42) + cs.add([ + Integer("top_k", (1, 20), default=4), + #Integer("chunk_size", (200, 4000), default=2000), + #Integer("chunk_overlap", (0, 500), default=200), + Categorical("prompt_template", list(prompt_map.keys()), default="default"), + ]) + #cs.add(ForbiddenGreaterThanRelation(cs["chunk_overlap"], cs["chunk_size"])) + scenario = Scenario( + cs, + deterministic=True, + n_trials=n_trials, + walltime_limit=7200, + seed=42, + output_directory=out / "smac_output", + ) + state: Dict[str, Any] = {"trial_count": 0, "best_cost": 1.0, "convergence": [], "history": []} + + smac = HyperparameterOptimizationFacade( + scenario=scenario, + target_function=lambda config, seed=0: objective(config, questions, dataset_name, state, answers_dir,prompt_map) + ) + t0 = time.time() + incumbent = smac.optimize() + if isinstance(incumbent, list): + best_cfg = incumbent[0].get_dictionary() + else: + best_cfg = incumbent.get_dictionary() + importance = param_importance(smac, out) + total_time = time.time() - t0 + #best_cfg = incumbent.get_dictionary() + results_path = out / "optimization_results.json" + results_path.write_text(json.dumps({ + "best_config": best_cfg, + "run_history": state["history"], + "convergence": state["convergence"], + "parameter_importance": importance, + "total_time_s": round(total_time, 2), + "dataset": dataset_name, + "metric": "bertscore_golden", + "answers_dir":str(answers_dir), + }, indent=2, default = str)) + + return best_cfg, state["history"], importance + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("task", nargs='?', default="learn2rag.optimization") + parser.add_argument("--dataset", type=str, default="WikiEval") + parser.add_argument("--max_questions", type=int, default=50) + parser.add_argument("--n_trials", type=int, default=10) + parser.add_argument("--logging-config", type=str) + parser.add_argument("--registry", type=str, default="registry.json") + parser.add_argument("--output_dir", type=str, default="optimization_results_baseline") + args, _ = parser.parse_known_args() + + final_output_dir = pathlib.Path(args.output_dir) + + env_out = os.environ.get("PIPELINE_OPT_CONFIG") + if not final_output_dir.exists() and env_out: + final_output_dir = pathlib.Path(env_out).parent + + incumbent, history, importance = run(args.dataset, args.max_questions, args.n_trials, final_output_dir, args.registry) + + # incumbent, history, importance = run( + # args.dataset, args.max_questions, args.n_trials, args.output_dir, + # ) + + best = min(history, key=lambda x: x["cost"]) + print(f"\nBest config: {dict(incumbent)}") + print(f"BERTScore (golden): {best['avg_bertscore_golden']:.4f}") + if importance: + print(f"\nParameter importance:") + for i, p in enumerate(importance["ranking"], 1): + print(f" {i}. {p}: {importance['individual'][p]:.4f}") \ No newline at end of file diff --git a/learn2rag/optimization/contrastive_optimization.py b/learn2rag/optimization/contrastive_optimization.py new file mode 100644 index 0000000..91a5d50 --- /dev/null +++ b/learn2rag/optimization/contrastive_optimization.py @@ -0,0 +1,333 @@ +""" +RAG Pipeline Optimization with contrastive evaluation. +""" + +import argparse +import json +import logging +import pathlib +import time +import copy + +import numpy as np +from bert_score import score as bert_score +from ConfigSpace import ConfigurationSpace, Integer, Categorical, ForbiddenGreaterThanRelation +from smac import HyperparameterOptimizationFacade, Scenario + +from learn2rag.evaluation.tools import read_dataset_qa +from learn2rag.pipeline.config import opt_config +import learn2rag.pipeline.search +import learn2rag.pipeline.generate + + +DATASET_CONFIG = { + "WikiEval": { + "subdirectory": "", + "split": "train", + "question_field": "question", + "answer_field": "answer", + "id_field": "id", + }, + "rag-mini-bioasq": { + "subdirectory": "question-answer-passages", + "split": "test", + "question_field": "question", + "answer_field": "answer", + "id_field": "id", + }, + "hotpot_qa": { # Not being used + "subdirectory": "distractor", + "split": "validation", + "question_field": "question", + "answer_field": "answer", + "id_field": "id", + }, + "repliqa": { # Not being used + "subdirectory": "repliqa_4", + "split": None, + "question_field": "question", + "answer_field": "long_answer", + "id_field": "question_id", + }, +} + +PROMPT_MAP = { + "default": ( + "# Role and Objective\nYou will act as a smart AI chatbot that answers " + "questions only by using the content from the provided information list.\n\n" + "# Instructions\n- Respond in the language of the question.\n" + "- Answer clear and concise.\n- Only use the provided information.\n" + "- NEVER use your general knowledge.\n\n" + "# Information:\n{context}" + ), + "concise": ( + "Answer the question using ONLY the provided information. " + "Be concise and direct. If the information does not contain the answer, say so.\n\n" + "Information:\n{context}" + ), + "detailed": ( + "You are a knowledgeable assistant. Using ONLY the provided information below, " + "answer the question thoroughly. Cite your sources. " + "If the information is insufficient, state that clearly.\n\n" + "Information:\n{context}" + ), +} + + +def load_false_answers(contrastive_file): + with open(contrastive_file) as f: + data = json.load(f) + return {item["question"]: item["false_answer"] for item in data} + + +def run_pipeline(question, user_config, working_config): + t0 = time.time() + docs = learn2rag.pipeline.search.search(question, user_config, working_config) + search_time = time.time() - t0 + + t0 = time.time() + answer = learn2rag.pipeline.generate.generate(question, docs, working_config) + gen_time = time.time() - t0 + + doc_list = docs.points if hasattr(docs, "points") else docs + context = "" + if doc_list: + context = "\n\n".join([ + f"Source: {d.payload.get('path', 'unknown')}\nContent: {d.payload.get('content', '')}" + for d in doc_list + ]) + return answer, context[:3000], search_time, gen_time + + +def objective(config, seed, questions, dataset_name, false_map, state, answers_dir): + state["trial_count"] += 1 + tid = state["trial_count"] + cfg = dict(config) + logging.info(f"Trial {tid}: {cfg}") + + wcfg = copy.deepcopy(opt_config) + wcfg["top_k"] = cfg["top_k"] + wcfg["chunk_size"] = cfg["chunk_size"] + wcfg["chunk_overlap"] = cfg["chunk_overlap"] + wcfg["prompt"] = PROMPT_MAP[cfg["prompt_template"]] + ucfg = {"file_path": None, "collection_name": dataset_name, + "imported_documents_file_path": None, "llm": None} + + predictions, goldens, falses = [], [], [] + qa_pairs = [] + t_start = time.time() + t_search, t_gen = 0.0, 0.0 + + for idx, q in enumerate(questions): + if not q["question"]: + continue + false_ans = false_map.get(q["question"], "") + if not false_ans: + continue + try: + answer, context, st, gt = run_pipeline(q["question"], ucfg, wcfg) + t_search += st + t_gen += gt + predictions.append(answer) + goldens.append(q["ground_truth"]) + falses.append(false_ans) + qa_pairs.append({ + "id": q["id"], + "question": q["question"], + "golden_answer": q["ground_truth"], + "generated_answer": answer, + "retrieved_context": context, + }) + except Exception as e: + logging.warning(f"Trial {tid}, q{idx} failed: {e}") + predictions.append("") + goldens.append(q["ground_truth"]) + falses.append(false_ans) + qa_pairs.append({ + "id": q["id"], + "question": q["question"], + "golden_answer": q["ground_truth"], + "generated_answer": "", + "retrieved_context": "", + }) + + if not predictions: + return 1.0 + + t_score = time.time() + _, _, F1_gold = bert_score(predictions, goldens, lang="en", verbose=False, rescale_with_baseline=True) + _, _, F1_false = bert_score(predictions, falses, lang="en", verbose=False, rescale_with_baseline=True) + scoring_time = time.time() - t_score + + bert_gold = [max(0.0, f.item()) for f in F1_gold] + bert_false = [max(0.0, f.item()) for f in F1_false] + + ratios = [] + # objective function + for bg, bf in zip(bert_gold, bert_false): + denom = bg + bf + ratios.append(bg / denom if denom > 0 else 0.5) + + avg_ratio = np.mean(ratios) + avg_bert_gold = np.mean(bert_gold) + avg_bert_false = np.mean(bert_false) + cost = 1.0 - avg_ratio + total_time = time.time() - t_start + + trial_answers = { + "trial_id": tid, + "config": cfg, + "cost": float(cost), + "avg_ratio": float(avg_ratio), + "qa_pairs": qa_pairs, + } + answers_file = answers_dir / f"trial_{tid}_answers.json" + with open(answers_file, "w") as f: + json.dump(trial_answers, f, indent=2, default=str) + + state["best_cost"] = min(state["best_cost"], cost) + state["convergence"].append({"trial": tid, "cost": float(cost), "best_cost": float(state["best_cost"])}) + state["history"].append({ + "trial_id": tid, "config": cfg, + "avg_ratio": float(avg_ratio), + "avg_bertscore_golden": float(avg_bert_gold), + "avg_bertscore_false": float(avg_bert_false), + "cost": float(cost), "time_s": round(total_time, 2), + "search_s": round(t_search, 2), "gen_s": round(t_gen, 2), + "scoring_s": round(scoring_time, 2), + }) + + logging.info( + f"Trial {tid}: ratio={avg_ratio:.4f} [gold={avg_bert_gold:.4f} false={avg_bert_false:.4f}] " + f"time={total_time:.1f}s (search={t_search:.1f} gen={t_gen:.1f} score={scoring_time:.1f})" + ) + return float(cost) + + +def param_importance(smac, output_path): + params = list(smac.scenario.configspace.keys()) + configs, costs = [], [] + for key, val in smac.runhistory.items(): + configs.append(dict(smac.runhistory.get_config(key.config_id))) + costs.append(val.cost) + if len(configs) < 3: + return {} + + raw = {} + for p in params: + groups = {} + for c, cost in zip(configs, np.array(costs)): + groups.setdefault(str(c[p]), []).append(cost) + means = [np.mean(g) for g in groups.values()] + raw[p] = float(np.var(means)) if len(means) > 1 else 0.0 + + total = sum(raw.values()) + imp = {p: round(v / total, 4) for p, v in raw.items()} if total > 0 else raw + ranking = sorted(imp, key=imp.get, reverse=True) + result = {"method": "variance_based", "ranking": ranking, "individual": imp} + with open(output_path / "parameter_importance.json", "w") as f: + json.dump(result, f, indent=2) + return result + + +def run(dataset_name, max_questions, n_trials, output_dir, contrastive_dir): + if dataset_name not in DATASET_CONFIG: + raise ValueError( + f"Unknown dataset: {dataset_name}. " + f"Available: {list(DATASET_CONFIG.keys())}" + ) + + dcfg = DATASET_CONFIG[dataset_name] + contrastive_file = pathlib.Path(contrastive_dir) / f"contrastive_answers_{dataset_name}.json" + out = pathlib.Path(output_dir) / dataset_name + out.mkdir(parents=True, exist_ok=True) + + answers_dir = out / "trial_answers" + answers_dir.mkdir(parents=True, exist_ok=True) + + false_map = load_false_answers(contrastive_file) + logging.info(f"Loaded {len(false_map)} false answers from {contrastive_file}") + + qa = read_dataset_qa(dataset_name, dcfg["subdirectory"], dcfg["split"]) + if max_questions: + qa = qa.select(range(min(max_questions, len(qa)))) + + questions = [] + for i, r in enumerate(qa): + questions.append({ + "question": r.get(dcfg["question_field"], ""), + "ground_truth": r.get(dcfg["answer_field"], ""), + "id": r.get(dcfg["id_field"], str(i)), + }) + logging.info(f"Loaded {len(questions)} questions from {dataset_name}") + + cs = ConfigurationSpace(seed=42) + cs.add([Integer("top_k", (1, 20), default=4), + Integer("chunk_size", (200, 4000), default=2000), + Integer("chunk_overlap", (0, 500), default=200), + Categorical("prompt_template", ["default", "concise", "detailed"], default="default")]) + cs.add(ForbiddenGreaterThanRelation(cs["chunk_overlap"], cs["chunk_size"])) + + scenario = Scenario(configspace=cs, deterministic=True, n_trials=n_trials, + walltime_limit=7200, seed=42, output_directory=out / "smac_output") + + state = {"trial_count": 0, "best_cost": 1.0, "convergence": [], "history": []} + smac = HyperparameterOptimizationFacade( + scenario=scenario, + target_function=lambda config, seed=0: objective( + config, seed, questions, dataset_name, false_map, state, answers_dir + ), + ) + + t0 = time.time() + incumbent = smac.optimize() + total_time = time.time() - t0 + + importance = param_importance(smac, out) + + with open(out / "optimization_results.json", "w") as f: + json.dump({ + "best_config": dict(incumbent), + "run_history": state["history"], + "convergence": state["convergence"], + "parameter_importance": importance, + "total_time_s": round(total_time, 2), + "dataset": dataset_name, + "metric": "bertscore_ratio(golden, false)", + "answers_dir": str(answers_dir), + }, f, indent=2, default=str) + + logging.info(f"Done in {total_time:.0f}s") + logging.info(f"Trial answers saved to {answers_dir}") + return incumbent, state["history"], importance + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--dataset", type=str, default="WikiEval", + choices=list(DATASET_CONFIG.keys())) + parser.add_argument("--max_questions", type=int, default=50) + parser.add_argument("--n_trials", type=int, default=10) + parser.add_argument("--output_dir", type=str, default="optimization_results_contrastive") + parser.add_argument("--contrastive_dir", type=str, default="contrastive_answers") + args = parser.parse_args() + + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(levelname)s %(message)s", + force=True, + ) + incumbent, history, importance = run( + args.dataset, args.max_questions, args.n_trials, + args.output_dir, args.contrastive_dir, + ) + + best = min(history, key=lambda x: x["cost"]) + print(f"\nBest config: {dict(incumbent)}") + print(f"Ratio score: {best['avg_ratio']:.4f}") + print(f" BERTScore vs golden: {best['avg_bertscore_golden']:.4f}") + print(f" BERTScore vs false: {best['avg_bertscore_false']:.4f}") + if importance: + print(f"\nParameter importance:") + for i, p in enumerate(importance["ranking"], 1): + print(f" {i}. {p}: {importance['individual'][p]:.4f}") \ No newline at end of file diff --git a/learn2rag/optimization/generate_contrastive_answers.py b/learn2rag/optimization/generate_contrastive_answers.py new file mode 100644 index 0000000..feb233f --- /dev/null +++ b/learn2rag/optimization/generate_contrastive_answers.py @@ -0,0 +1,125 @@ +""" +Generate contrastive (false) answers for RAG evaluation. +""" + +import argparse +import json +import logging +import pathlib +import time + +from learn2rag.pipeline.llm import llm as learn2rag_llm +from learn2rag.evaluation.tools import read_dataset_qa +from langchain_core.messages import HumanMessage + + +DATASET_CONFIG = { + "WikiEval": { + "subdirectory": "", + "split": "train", + "question_field": "question", + "answer_field": "answer", + "id_field": "id", + }, + "rag-mini-bioasq": { + "subdirectory": "question-answer-passages", + "split": "test", + "question_field": "question", + "answer_field": "answer", + "id_field": "id", + }, + "hotpot_qa": { # Not being used + "subdirectory": "distractor", + "split": "validation", + "question_field": "question", + "answer_field": "answer", + "id_field": "id", + }, + "repliqa": { # Not being used + "subdirectory": "repliqa_4", + "split": None, + "question_field": "question", + "answer_field": "long_answer", + "id_field": "question_id", + }, +} + + +def generate(prompt): + response = learn2rag_llm.invoke([HumanMessage(content=prompt)]) + return response.content.strip() + + +def generate_false_answer(question): + return generate( + f"Answer the given question in an incorrect manner.\n\n" + f"question: {question}" + ) + +def run(dataset_name, max_questions=50, output_dir="contrastive_answers"): + if dataset_name not in DATASET_CONFIG: + raise ValueError( + f"Unknown dataset: {dataset_name}. " + f"Available: {list(DATASET_CONFIG.keys())}" + ) + + cfg = DATASET_CONFIG[dataset_name] + output_path = pathlib.Path(output_dir) + output_path.mkdir(parents=True, exist_ok=True) + + qa = read_dataset_qa(dataset_name, cfg["subdirectory"], cfg["split"]) + if max_questions: + qa = qa.select(range(min(max_questions, len(qa)))) + + logging.info(f"Dataset: {dataset_name}, questions: {len(qa)}") + + results = [] + total_start = time.time() + + for idx, item in enumerate(qa): + question = item.get(cfg["question_field"], "") + golden = item.get(cfg["answer_field"], "") + qid = item.get(cfg["id_field"], str(idx)) + + if not question: + continue + + print(f"[{idx+1}/{len(qa)}] {question[:60]}...") + + t0 = time.time() + false_answer = generate_false_answer(question) + elapsed = time.time() - t0 + + results.append({ + "id": qid, + "question": question, + "golden_answer": golden, + "false_answer": false_answer, + "generation_time_s": round(elapsed, 2), + }) + + total_time = time.time() - total_start + + out_file = output_path / f"contrastive_answers_{dataset_name}.json" + with open(out_file, "w") as f: + json.dump(results, f, indent=2) + + print(f"\nGenerated false answers for {len(results)} questions in {total_time:.0f}s") + print(f"Saved to {out_file}") + return results + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--dataset", type=str, default="WikiEval", + choices=list(DATASET_CONFIG.keys())) + parser.add_argument("--max_questions", type=int, default=50) + parser.add_argument("--output_dir", type=str, default="contrastive_answers") + args = parser.parse_args() + + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(levelname)s %(message)s", + force=True, + ) + run(args.dataset, args.max_questions, args.output_dir) \ No newline at end of file diff --git a/learn2rag/optimization/readme.md b/learn2rag/optimization/readme.md new file mode 100644 index 0000000..e69de29 diff --git a/learn2rag/optimization/registry.json b/learn2rag/optimization/registry.json new file mode 100644 index 0000000..3377340 --- /dev/null +++ b/learn2rag/optimization/registry.json @@ -0,0 +1,19 @@ +{ + "datasets": { + "WikiEval": { + "subdirectory": "", "split": "train", + "fields": {"q": "question", "a": "answer", "id": "id"}, + "path": "./CSV/WikiEval_dataset.csv" + } + }, + "prompts": { + "default": + "# Role and Objective\nYou will act as a smart AI chatbot that answers questions only by using the content from the provided information list.\n\n # Instructions\n- Respond in the language of the question.\n - Answer clear and concise.\n- Only use the provided information.\n - NEVER use your general knowledge.\n\n # Information:\n{context}" + , + "concise": + "Answer the question using ONLY the provided information. Be concise and direct. If the information does not contain the answer, say so.\n\n Information:\n{context}" + , + "detailed": + "You are a knowledgeable assistant. Using ONLY the provided information below, answer the question thoroughly. Cite your sources. If the information is insufficient, state that clearly.\n\n Information:\n{context}" + } +} \ No newline at end of file diff --git a/learn2rag/optimization/retrieval_optimization.py b/learn2rag/optimization/retrieval_optimization.py new file mode 100644 index 0000000..586d917 --- /dev/null +++ b/learn2rag/optimization/retrieval_optimization.py @@ -0,0 +1,703 @@ +""" +RAG Retrieval Optimization. +""" +import os +os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" +os.environ["CUDA_VISIBLE_DEVICES"]="0" #Select GPU number 0 + +import argparse +import datetime +import json +import logging +import pathlib +import time +import copy +import os +import asyncio +import subprocess +import sys +from typing import Dict, Any, List, Union, Tuple, cast + +import numpy as np +# from bert_score import score as bert_score # type: ignore[import-not-found] +from ConfigSpace import ( + ConfigurationSpace, + Integer, + Categorical, + ForbiddenGreaterThanRelation, + Configuration, + ForbiddenAndConjunction, + ForbiddenEqualsClause, +) +from smac import HyperparameterOptimizationFacade, Scenario + +from learn2rag.evaluation.tools import read_dataset_qa +from learn2rag.pipeline.config import opt_config +import learn2rag.pipeline.search +# import learn2rag.pipeline.generate + + + + +def load_registry(path: str = "registry.json") -> dict[str, Any]: + p = pathlib.Path(path) + if not p.is_file(): + logging.error("registry file not found") + with p.open() as f: + return cast(dict[str, Any], json.load(f)) + + +def _load_existing_trial_answers(answers_dir: pathlib.Path) -> list[dict[str, Any]]: + if not answers_dir.exists(): + return [] + trials: list[dict[str, Any]] = [] + for p in sorted(answers_dir.glob("trial_*_answers.json")): + try: + data = json.loads(p.read_text()) + if isinstance(data, dict) and "trial_id" in data: + trials.append(data) + except Exception as e: + logging.warning(f"Could not read {p}: {e}") + return trials + + +def _restore_state_from_existing(out: pathlib.Path, answers_dir: pathlib.Path) -> Dict[str, Any]: + state: Dict[str, Any] = {"trial_count": 0, "best_cost": 1.0, "convergence": [], "history": []} + + trial_answers = _load_existing_trial_answers(answers_dir) + answers_by_id = { + int(t.get("trial_id", 0)): t + for t in trial_answers + if isinstance(t, dict) and t.get("trial_id") is not None + } + + results_path = out / "optimization_results.json" + if results_path.exists(): + try: + results = json.loads(results_path.read_text()) + history = results.get("run_history", []) + convergence = results.get("convergence", []) + if isinstance(history, list) and history: + merged_history: dict[int, dict[str, Any]] = { + int(h.get("trial_id", 0)): dict(h) + for h in history + if isinstance(h, dict) and h.get("trial_id") is not None + } + + for tid, trial in answers_by_id.items(): + previous = merged_history.get(tid, {}) + merged_history[tid] = { + "trial_id": tid, + "config": trial.get("config", previous.get("config", {})), + "recall": float(trial.get("recall") or previous.get("recall") or 0.0), + "avg_t_search": float(trial.get("avg_t_search") or previous.get("avg_t_search") or 0.0), + "cost": float(trial.get("cost") or previous.get("cost") or 1.0), + "time_s": previous.get("time_s"), + "search_s": previous.get("search_s"), + "scoring_s": previous.get("scoring_s"), + } + + merged_list = [merged_history[tid] for tid in sorted(merged_history)] + best_cost = 1.0 + rebuilt_convergence: list[dict[str, Any]] = [] + for entry in merged_list: + best_cost = min(best_cost, float(entry.get("cost", 1.0))) + rebuilt_convergence.append({ + "trial": int(entry.get("trial_id", 0)), + "cost": float(entry.get("cost", 1.0)), + "best_cost": best_cost, + }) + + state["history"] = merged_list + state["convergence"] = rebuilt_convergence if not isinstance(convergence, list) or len(rebuilt_convergence) != len(convergence) else convergence + state["trial_count"] = max(int(h.get("trial_id", 0)) for h in merged_list) + state["best_cost"] = best_cost + return state + except Exception as e: + logging.warning(f"Could not read {results_path}: {e}") + + trials = trial_answers + if not trials: + return state + + best_cost = 1.0 + new_history: list[dict[str, Any]] = [] + new_convergence: list[dict[str, Any]] = [] + for trial in sorted(trials, key=lambda t: int(t.get("trial_id", 0))): + tid = int(trial.get("trial_id", 0)) + cost = float(trial.get("cost", 1.0)) + best_cost = min(best_cost, cost) + new_history.append({ + "trial_id": tid, + "config": trial.get("config", {}), + "recall": float(trial.get("recall", 0.0)), + "avg_t_search": float(trial.get("avg_t_search", 0.0)), + "cost": cost, + "time_s": None, + "search_s": None, + "scoring_s": None, + }) + new_convergence.append({"trial": tid, "cost": cost, "best_cost": best_cost}) + + state["history"] = new_history + state["convergence"] = new_convergence + state["trial_count"] = max(int(t.get("trial_id", 0)) for t in trials) + state["best_cost"] = best_cost + return state + + +def _load_existing_importance(out: pathlib.Path) -> Dict[str, Any]: + path = out / "parameter_importance.json" + if not path.exists(): + return {} + try: + data = json.loads(path.read_text()) + return data if isinstance(data, dict) else {} + except Exception as e: + logging.warning(f"Could not read {path}: {e}") + return {} + + +def _find_latest_optimization_file(smac_output_dir: pathlib.Path) -> Union[pathlib.Path, None]: + if not smac_output_dir.exists(): + return None + candidates = list(smac_output_dir.rglob("optimization.json")) + if not candidates: + return None + return max(candidates, key=lambda p: p.stat().st_mtime) + + +def _parse_last_update(value: Any) -> Union[float, None]: + if isinstance(value, (int, float)): + return float(value) + if isinstance(value, str): + try: + # Handle ISO values like 2026-06-11T12:34:56.123456+00:00 or trailing Z. + dt = datetime.datetime.fromisoformat(value.replace("Z", "+00:00")) + return dt.timestamp() + except ValueError: + return None + return None + + +def _last_update_age_seconds(smac_output_dir: pathlib.Path) -> Union[float, None]: + optimization_file = _find_latest_optimization_file(smac_output_dir) + if optimization_file is None: + return None + try: + data = json.loads(optimization_file.read_text()) + except Exception as e: + logging.warning(f"Could not read {optimization_file}: {e}") + return None + ts = _parse_last_update(data.get("last_update")) if isinstance(data, dict) else None + if ts is None: + return None + return max(0.0, time.time() - ts) + + +def _run_search_heartbeat_age_seconds(heartbeat_file: pathlib.Path) -> Union[float, None]: + if not heartbeat_file.exists(): + return None + try: + return max(0.0, time.time() - heartbeat_file.stat().st_mtime) + except OSError as e: + logging.warning(f"Could not stat {heartbeat_file}: {e}") + return None + + +def _touch_run_search_heartbeat() -> None: + heartbeat_path = os.environ.get("L2R_RUN_SEARCH_HEARTBEAT_FILE") + if not heartbeat_path: + return + try: + p = pathlib.Path(heartbeat_path) + p.parent.mkdir(parents=True, exist_ok=True) + p.touch() + except Exception as e: + logging.warning(f"Could not update run_search heartbeat at {heartbeat_path}: {e}") + + +def _build_worker_command(args: argparse.Namespace, final_output_dir: pathlib.Path) -> List[str]: + cmd = [ + sys.executable, + "-m", + "learn2rag.optimization.retrieval_optimization", + "--dataset", args.dataset, + "--max_questions", str(args.max_questions), + "--n_trials", str(args.n_trials), + "--registry", args.registry, + "--output_dir", str(final_output_dir), + "--resume", + ] + if args.n_trials_is_total or args.resume: + cmd.append("--n_trials_is_total") + if args.logging_config: + cmd.extend(["--logging-config", args.logging_config]) + return cmd + + +def run_with_watchdog(args: argparse.Namespace, final_output_dir: pathlib.Path) -> int: + stale_after_s = max(1, args.watchdog_stale_minutes * 60) + run_search_stale_after_s = max(1, args.watchdog_run_search_stale_minutes * 60) + restart_wait_s = max(1, args.watchdog_restart_delay_minutes * 60) + poll_s = max(5, args.watchdog_poll_seconds) + + dataset_out = final_output_dir / args.dataset + smac_output_dir = dataset_out / "smac_output" + run_search_heartbeat_file = dataset_out / "run_search_heartbeat.txt" + worker_cmd = _build_worker_command(args, final_output_dir) + + restart_count = 0 + while True: + logging.info(f"Starting optimization worker (restart #{restart_count})") + try: + run_search_heartbeat_file.unlink(missing_ok=True) + except OSError as e: + logging.warning(f"Could not reset heartbeat file {run_search_heartbeat_file}: {e}") + + worker_env = os.environ.copy() + worker_env["L2R_RUN_SEARCH_HEARTBEAT_FILE"] = str(run_search_heartbeat_file) + proc = subprocess.Popen(worker_cmd, env=worker_env) + stale_detected = False + + while proc.poll() is None: + time.sleep(poll_s) + age = _last_update_age_seconds(smac_output_dir) + if age is not None and age > stale_after_s: + stale_detected = True + logging.warning( + f"Detected stale optimization.json update (age={age:.0f}s > {stale_after_s}s). " + "Terminating worker for restart." + ) + proc.terminate() + try: + proc.wait(timeout=30) + except subprocess.TimeoutExpired: + proc.kill() + proc.wait() + break + + run_search_age = _run_search_heartbeat_age_seconds(run_search_heartbeat_file) + if run_search_age is not None and run_search_age > run_search_stale_after_s: + stale_detected = True + logging.warning( + f"Detected stale run_search heartbeat (age={run_search_age:.0f}s > {run_search_stale_after_s}s). " + "Terminating worker for restart." + ) + proc.terminate() + try: + proc.wait(timeout=30) + except subprocess.TimeoutExpired: + proc.kill() + proc.wait() + break + + if not stale_detected: + return_code = proc.returncode if proc.returncode is not None else 1 + if return_code == 0: + logging.info("Optimization worker finished successfully.") + else: + logging.error(f"Optimization worker exited with return code {return_code}.") + return return_code + + restart_count += 1 + logging.info(f"Sleeping {restart_wait_s}s before resuming optimization.") + time.sleep(restart_wait_s) + +def run_search(question: str, user_config: Dict[str, Any], working_config: Dict[str, Any]) -> Tuple[List[Any], float]: + _touch_run_search_heartbeat() + t0 = time.time() + docs = asyncio.run(learn2rag.pipeline.search.search_authorized(question, user="anonymous", request_id=None, user_config=user_config, opt_config=working_config)) + search_time = time.time() - t0 + source_list = [point.payload['source'] for point in docs if point.payload is not None and "source" in point.payload] + return source_list, search_time + + +def recall(search_results: list[list[Any]], labels: list[Any]) -> float: + count = 0 + top_k = opt_config["top_k"] + for q in range(len(search_results)): + label = str(labels[q]) + hits = [str(h) for h in search_results[q]] + # print('label ', label, ' hits: ', hits) + if label in hits[:top_k]: + count += 1 + return count / len(labels) if labels else 0.0 + + +#I removed seed because there are no use for it +# removed dataset_name because it just use in user config and now we inject it +def objective(config: Configuration, + questions: List[Dict[str, Any]], + state: Dict[str, Any], + answers_dir: pathlib.Path +) -> float: + state["trial_count"] += 1 + tid = state["trial_count"] + cfg = dict(config) + logging.info(f"Trial {tid}: {cfg}") + + if cfg["rewrite"] == "False" and cfg["rewrite_mode"] in {"keywords", "subqueries", "subqueries_keywords"}: + logging.warning(f"Skip invalid cfg: {cfg}") + return 1.0 + + if cfg["reranking"] == "False" and cfg["reranking_mode"] in {"reranking_with_flagreranker", "reranking_with_sentence_transformers", "reranking_with_colbert"}: + logging.warning(f"Skip invalid cfg: {cfg}") + return 1.0 + + if cfg["search_mode"] in {"dense", "sparse"} and cfg["fusion_mode"] in {"DBSF", "RRF"}: + logging.warning(f"Skip invalid cfg: {cfg}") + return 1.0 + + working_cfg = copy.deepcopy(opt_config) + working_cfg.update({ + "chunk_size": cfg["chunk_size"], + "chunk_overlap": cfg["chunk_overlap"], + "search_mode": cfg["search_mode"], + "reranking_mode": cfg["reranking_mode"], + "rewrite_mode": cfg["rewrite_mode"], + "fusion_mode": cfg["fusion_mode"], + "reranking": cfg["reranking"], + "rewrite": cfg["rewrite"] + }) + + ucfg = { + "file_path": None, + "collection_name": f"CSC-CS_{cfg['chunk_size']}-CO_{cfg['chunk_overlap']}", + "imported_documents_file_path": None, + "llm": None, + } + env_user_cfg = os.environ.get("PIPELINE_USER_CONFIG") + if env_user_cfg and pathlib.Path(env_user_cfg).exists(): + ucfg.update(json.loads(pathlib.Path(env_user_cfg).read_text())) + + predictions: list[list[Any]] = [] + goldens: list[Any] = [] + qa_pairs: list[dict[str, Any]] = [] + + t_start = time.time() + t_search = 0.0 + + for q in questions: + # Preserve the original “skip empty question” guard + if not q.get("question"): + continue + try: + source_list, t_s = run_search(q["question"], ucfg, working_cfg) + t_search += t_s + predictions.append(source_list) + goldens.append(q["ground_truth"]) + qa_pairs.append({**q, "retrieved_sources": source_list}) + except Exception as e: + # Same behaviour as the old version: record a blank answer. TODO : check if we need this + logging.warning(f"Trial {tid}, q{q.get('id','?')} failed: {e}") + predictions.append([]) + goldens.append(q["ground_truth"]) + qa_pairs.append({**q, "retrieved_sources": ""}) + + if not predictions: + return 1.0 + + t_score = time.time() + recall_score = recall(predictions, goldens) + scoring_time = time.time() - t_score + total_time = time.time() - t_start + + # objective function + w_recall = 0.5 + w_time = 0.5 + t_search_per_sample_upper = 50 + max_time_s = t_search_per_sample_upper*len(predictions) + time_cost = max(0.0, 1.0 - (t_search / max_time_s)) + cost = 1.0 - w_recall*recall_score - w_time*time_cost + avg_t_search = t_search/len(predictions) + + trial_answers = { + "trial_id": tid, + "config": cfg, + "cost": float(cost), + "recall": float(recall_score), + "avg_t_search": float(avg_t_search), + "w_recall": w_recall, + "w_time": w_time, + "top_k": opt_config["top_k"], + "qa_pairs": qa_pairs + } + answers_file = answers_dir / f"trial_{tid}_answers.json" + while answers_file.exists(): + # Keep IDs monotonic when resuming from a partially persisted run. + tid += 1 + state["trial_count"] = tid + trial_answers["trial_id"] = tid + answers_file = answers_dir / f"trial_{tid}_answers.json" + with open(answers_file, "w") as f: + json.dump(trial_answers, f, indent=2, default=str) + + state["best_cost"] = min(state["best_cost"], cost) + state["convergence"].append({"trial": tid, "cost": float(cost), "best_cost": float(state["best_cost"])}) + state["history"].append({ + "trial_id": tid, "config": cfg, + "recall": float(recall_score), + "avg_t_search": float(avg_t_search), + "cost": float(cost), + "time_s": round(total_time, 2), + "search_s": round(t_search, 2), + "scoring_s": round(scoring_time, 2), + }) + + logging.info( + f"Trial {tid}: recall={recall_score:.4f} avg_t_search={avg_t_search: .4f} cost={cost:.4f} " + f"time={total_time:.1f}s (search={t_search:.1f} score={scoring_time:.1f})" + ) + return float(cost) + + +def param_importance(smac: HyperparameterOptimizationFacade, output_path: pathlib.Path) -> Dict[str, Any]: + params = list(smac.scenario.configspace.keys()) + configs, costs = [], [] + for key, val in smac.runhistory.items(): + configs.append(dict(smac.runhistory.get_config(key.config_id))) + costs.append(val.cost) + if len(configs) < 3: + return {} + + raw: dict[str, float] = {} + for p in params: + groups: dict[str, list[float]] = {} + for c, cost in zip(configs, np.array(costs)): + groups.setdefault(str(c[p]), []).append(float(cost)) + means = [np.mean(g) for g in groups.values()] + raw[p] = float(np.var(means)) if len(means) > 1 else 0.0 + + total = sum(raw.values()) + imp = {p: round(v / total, 4) for p, v in raw.items()} if total > 0 else raw + ranking = sorted(imp, key=imp.get, reverse=True) # type: ignore[arg-type] + result = {"method": "variance_based", "ranking": ranking, "individual": imp} + with open(output_path / "parameter_importance.json", "w") as f: + json.dump(result, f, indent=2) + return result + + +def run( + dataset_name: str, + max_questions: int, + n_trials: int, + output_dir: Union[str, pathlib.Path], + registry_path: str, + resume: bool = False, + n_trials_is_total: bool = True, +) -> Tuple[ + Dict[str, Any], List[Any], Dict[str, Any]]: + registry = load_registry(registry_path) + datasets = registry["datasets"] + if dataset_name not in datasets: + raise ValueError(f"Unknown dataset: {dataset_name}. Available: {list(datasets.keys())}") + dcfg = datasets[dataset_name] + fields = dcfg["fields"] + + out = pathlib.Path(output_dir) / dataset_name + out.mkdir(parents=True, exist_ok=True) + answers_dir = out / "trial_answers" + answers_dir.mkdir(parents=True, exist_ok=True) + base_path = pathlib.Path(dcfg["path"]) + if base_path.suffix.lower() == '.csv': + target_path = base_path + else: + target_path = base_path / 'source' / dcfg.get("subdirectory", "") + logging.info(f"target path is {target_path}") + qa = read_dataset_qa(target_path, split=dcfg["split"]) + if max_questions: + qa = qa.select(range(min(max_questions, len(qa)))) + + questions = [ + { + "question": r[fields["q"]], + "ground_truth": r[fields["a"]], + "id": r.get(fields["id"], str(i)), + } + for i, r in enumerate(qa) + ] + + prompt_map = registry["prompts"] + cs = ConfigurationSpace(seed=42) + cs.add([ + Categorical("chunk_size", [250, 1000, 2000], default=1000), + Categorical("chunk_overlap", [50, 200], default=50), + Categorical("search_mode", ["dense", "sparse", "dense_sparse", "dense_sparse_colbert"], default="dense"), + Categorical("reranking_mode", ["none", "reranking_with_flagreranker", "reranking_with_sentence_transformers", "reranking_with_colbert"], default="none"), + Categorical("rewrite_mode", ["none", "subqueries", "keywords", "subqueries_keywords"], default="none"), + Categorical("fusion_mode", ["none", "DBSF", "RRF"], default="none"), + Categorical("reranking", ["True", "False"], default="False"), + Categorical("rewrite", ["True", "False"], default="False"), + Integer("top_k", (1, 20), default=4), + Categorical("prompt_template", list(prompt_map.keys()), default="default"), + ]) + cs.add(ForbiddenGreaterThanRelation(cs["chunk_overlap"], cs["chunk_size"])) + + cs.add(ForbiddenAndConjunction( + ForbiddenEqualsClause(cs["chunk_size"], 250), + ForbiddenEqualsClause(cs["chunk_overlap"], 200), + )) + + for sm in ["dense", "sparse"]: + for fm in ["DBSF", "RRF"]: + cs.add(ForbiddenAndConjunction( + ForbiddenEqualsClause(cs["search_mode"], sm), + ForbiddenEqualsClause(cs["fusion_mode"], fm), + )) + + for sm in ["dense_sparse", "dense_sparse_colbert"]: + cs.add(ForbiddenAndConjunction( + ForbiddenEqualsClause(cs["search_mode"], sm), + ForbiddenEqualsClause(cs["fusion_mode"], "none"), + )) + + + for rrm in ["reranking_with_flagreranker", "reranking_with_sentence_transformers", "reranking_with_colbert"]: + cs.add(ForbiddenAndConjunction( + ForbiddenEqualsClause(cs["reranking"], "False"), + ForbiddenEqualsClause(cs["reranking_mode"], rrm), + )) + + cs.add(ForbiddenAndConjunction( + ForbiddenEqualsClause(cs["reranking"], "True"), + ForbiddenEqualsClause(cs["reranking_mode"], "none"), + )) + + for rwm in ["keywords", "subqueries", "subqueries_keywords"]: + cs.add(ForbiddenAndConjunction( + ForbiddenEqualsClause(cs["rewrite"], "False"), + ForbiddenEqualsClause(cs["rewrite_mode"], rwm), + )) + + cs.add(ForbiddenAndConjunction( + ForbiddenEqualsClause(cs["rewrite"], "True"), + ForbiddenEqualsClause(cs["rewrite_mode"], "none"), + )) + + + state: Dict[str, Any] + if resume: + state = _restore_state_from_existing(out, answers_dir) + else: + state = {"trial_count": 0, "best_cost": 1.0, "convergence": [], "history": []} + + already_done = int(state["trial_count"]) + target_trials = max(already_done, n_trials) if n_trials_is_total else already_done + n_trials + remaining_trials = max(0, target_trials - already_done) + + best_cfg: Dict[str, Any] = {} + importance: Dict[str, Any] = {} + total_time = 0.0 + + if remaining_trials > 0: + scenario = Scenario( + cs, + deterministic=True, + n_trials=target_trials, + walltime_limit=172800, #7200, + seed=42, + output_directory=out / "smac_output", + ) + + initial_design = HyperparameterOptimizationFacade.get_initial_design( + scenario=scenario, + additional_configs=[cs.get_default_configuration()], + ) + + smac = HyperparameterOptimizationFacade( + scenario=scenario, + target_function=lambda config, seed=0: objective(config, questions, state, answers_dir), + initial_design=initial_design + ) + + t0 = time.time() + incumbent = smac.optimize() + if isinstance(incumbent, list): + incumbent = incumbent[0] + importance = param_importance(smac, out) + total_time = time.time() - t0 + best_cfg = dict(incumbent) + else: + logging.info("No remaining trials to run. Returning existing results.") + if state["history"]: + best_cfg = min(state["history"], key=lambda h: h["cost"]).get("config", {}) + importance = _load_existing_importance(out) + + if not best_cfg and state["history"]: + best_cfg = min(state["history"], key=lambda h: h["cost"]).get("config", {}) + + best_trial = min( + (h for h in state["history"] if h.get("config") == best_cfg), + key=lambda h: h["cost"], + default=None, + ) + best_trial_id = best_trial["trial_id"] if best_trial else None + results_path = out / "optimization_results.json" + results_path.write_text(json.dumps({ + "best_config": best_cfg, + "best_trial_id": best_trial_id, + "run_history": state["history"], + "convergence": state["convergence"], + "parameter_importance": importance, + "total_time_s": round(total_time, 2), + "dataset": dataset_name, + "metric": "recall", + "answers_dir":str(answers_dir), + }, indent=2, default = str)) + + return best_cfg, state["history"], importance + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("task", nargs='?', default="learn2rag.optimization") + parser.add_argument("--dataset", type=str, default="WikiEval") + parser.add_argument("--max_questions", type=int, default=50) + parser.add_argument("--n_trials", type=int, default=10) + parser.add_argument("--logging-config", type=str) + parser.add_argument("--registry", type=str, default="registry.json") + parser.add_argument("--output_dir", type=str, default="optimization_results_baseline") + parser.add_argument("--resume", action="store_true") + parser.add_argument( + "--n_trials_is_total", + action="store_true", + help="Interpret --n_trials as the total desired trial count instead of additional trials.", + ) + parser.add_argument("--watchdog", action="store_true", help="Restart optimization if SMAC last_update is stale.") + parser.add_argument("--watchdog_stale_minutes", type=int, default=90) + parser.add_argument("--watchdog_run_search_stale_minutes", type=int, default=7) + parser.add_argument("--watchdog_restart_delay_minutes", type=int, default=5) + parser.add_argument("--watchdog_poll_seconds", type=int, default=60) + args, _ = parser.parse_known_args() + + final_output_dir = pathlib.Path(args.output_dir) + + env_out = os.environ.get("PIPELINE_OPT_CONFIG") + if not final_output_dir.exists() and env_out: + final_output_dir = pathlib.Path(env_out).parent + + if args.watchdog: + exit_code = run_with_watchdog(args, final_output_dir) + raise SystemExit(exit_code) + + incumbent, history, importance = run( + args.dataset, + args.max_questions, + args.n_trials, + final_output_dir, + args.registry, + resume=args.resume, + n_trials_is_total=(args.n_trials_is_total or args.resume), + ) + + # incumbent, history, importance = run( + # args.dataset, args.max_questions, args.n_trials, args.output_dir, + # ) + + best = min(history, key=lambda x: x["cost"]) if history else None + print(f"\nBest config: {dict(incumbent)}") + # print(f"BERTScore (golden): {best['avg_bertscore_golden']:.4f}") + if importance: + print(f"\nParameter importance:") + for i, p in enumerate(importance["ranking"], 1): + print(f" {i}. {p}: {importance['individual'][p]:.4f}") \ No newline at end of file diff --git a/learn2rag/optimization/self_judge_optimization.py b/learn2rag/optimization/self_judge_optimization.py new file mode 100644 index 0000000..93a9dc0 --- /dev/null +++ b/learn2rag/optimization/self_judge_optimization.py @@ -0,0 +1,340 @@ +""" +RAG pipeline optimization with self-judge using answer relevance only. +""" + +import argparse +import json +import logging +import pathlib +import time +import copy +import re + +import numpy as np +from sentence_transformers import SentenceTransformer +from ConfigSpace import ConfigurationSpace, Integer, Categorical, ForbiddenGreaterThanRelation +from smac import HyperparameterOptimizationFacade, Scenario + +from learn2rag.evaluation.tools import read_dataset_qa +from learn2rag.pipeline.config import opt_config +from learn2rag.pipeline.llm import llm as learn2rag_llm +from langchain_core.messages import HumanMessage +import learn2rag.pipeline.search +import learn2rag.pipeline.generate + + +DATASET_CONFIG = { + "WikiEval": { + "subdirectory": "", + "split": "train", + "question_field": "question", + "answer_field": "answer", + "id_field": "id", + }, + "rag-mini-bioasq": { + "subdirectory": "question-answer-passages", + "split": "test", + "question_field": "question", + "answer_field": "answer", + "id_field": "id", + }, + "hotpot_qa": { # Not being used + "subdirectory": "distractor", + "split": "validation", + "question_field": "question", + "answer_field": "answer", + "id_field": "id", + }, + "repliqa": { # Not being used + "subdirectory": "repliqa_4", + "split": None, + "question_field": "question", + "answer_field": "long_answer", + "id_field": "question_id", + }, +} + +PROMPT_MAP = { + "default": ( + "# Role and Objective\nYou will act as a smart AI chatbot that answers " + "questions only by using the content from the provided information list.\n\n" + "# Instructions\n- Respond in the language of the question.\n" + "- Answer clear and concise.\n- Only use the provided information.\n" + "- NEVER use your general knowledge.\n\n" + "# Information:\n{context}" + ), + "concise": ( + "Answer the question using ONLY the provided information. " + "Be concise and direct. If the information does not contain the answer, say so.\n\n" + "Information:\n{context}" + ), + "detailed": ( + "You are a knowledgeable assistant. Using ONLY the provided information below, " + "answer the question thoroughly. Cite your sources. " + "If the information is insufficient, state that clearly.\n\n" + "Information:\n{context}" + ), +} + +encoder = SentenceTransformer("all-MiniLM-L6-v2") + + +def call_llm(prompt, max_tokens=512): + response = learn2rag_llm.invoke([HumanMessage(content=prompt)], max_tokens=max_tokens) + return response.content.strip() + + +def answer_relevance(question, answer): + if not question or not answer: + return 0.0 + + generated_qs = [] + for _ in range(3): + text = call_llm( + f"Generate a question for the given answer.\n\nanswer: {answer}", + max_tokens=128, + ) + first_line = text.split("\n")[0].strip() + first_line = re.sub(r"^\d+[\.\)]\s*", "", first_line).strip() + if len(first_line) > 10: + generated_qs.append(first_line) + + if not generated_qs: + return 0.0 + + q_emb = encoder.encode([question])[0] + gen_embs = encoder.encode(generated_qs) + sims = [] + for ge in gen_embs: + cos = np.dot(q_emb, ge) / (np.linalg.norm(q_emb) * np.linalg.norm(ge) + 1e-8) + sims.append(max(0.0, float(cos))) + + return np.mean(sims) + + +def run_pipeline(question, user_config, working_config): + t0 = time.time() + docs = learn2rag.pipeline.search.search(question, user_config, working_config) + search_time = time.time() - t0 + + t0 = time.time() + answer = learn2rag.pipeline.generate.generate(question, docs, working_config) + gen_time = time.time() - t0 + + doc_list = docs.points if hasattr(docs, "points") else docs + context = "" + if doc_list: + context = "\n\n".join([ + f"Source: {d.payload.get('path', 'unknown')}\nContent: {d.payload.get('content', '')}" + for d in doc_list + ]) + return answer, context[:3000], search_time, gen_time + + +def objective(config, seed, questions, dataset_name, state, answers_dir): + state["trial_count"] += 1 + tid = state["trial_count"] + cfg = dict(config) + logging.info(f"Trial {tid}: {cfg}") + + wcfg = copy.deepcopy(opt_config) + wcfg["top_k"] = cfg["top_k"] + wcfg["chunk_size"] = cfg["chunk_size"] + wcfg["chunk_overlap"] = cfg["chunk_overlap"] + wcfg["prompt"] = PROMPT_MAP[cfg["prompt_template"]] + ucfg = {"file_path": None, "collection_name": dataset_name, + "imported_documents_file_path": None, "llm": None} + + rel_scores = [] + qa_pairs = [] + t_start = time.time() + t_search, t_gen, t_judge = 0.0, 0.0, 0.0 + + for idx, q in enumerate(questions): + if not q["question"]: + continue + try: + answer, context, st, gt = run_pipeline(q["question"], ucfg, wcfg) + t_search += st + t_gen += gt + + tj = time.time() + rel = answer_relevance(q["question"], answer) + rel_scores.append(rel) + t_judge += time.time() - tj + + qa_pairs.append({ + "id": q["id"], + "question": q["question"], + "ground_truth": q["ground_truth"], + "generated_answer": answer, + "retrieved_context": context, + "answer_relevance": float(rel), + }) + except Exception as e: + logging.warning(f"Trial {tid}, q{idx} failed: {e}") + rel_scores.append(0.0) + qa_pairs.append({ + "id": q["id"], + "question": q["question"], + "ground_truth": q["ground_truth"], + "generated_answer": "", + "retrieved_context": "", + "answer_relevance": 0.0, + }) + + if not rel_scores: + return 1.0 + + # Objective function + avg_rel = float(np.mean(rel_scores)) + cost = 1.0 - avg_rel + total_time = time.time() - t_start + + trial_data = { + "trial_id": tid, + "config": cfg, + "avg_answer_relevance": avg_rel, + "cost": float(cost), + "time_s": round(total_time, 2), + "qa_pairs": qa_pairs, + } + with open(answers_dir / f"trial_{tid}_answers.json", "w") as f: + json.dump(trial_data, f, indent=2) + + state["best_cost"] = min(state["best_cost"], cost) + state["convergence"].append({"trial": tid, "cost": float(cost), "best_cost": float(state["best_cost"])}) + state["history"].append({ + "trial_id": tid, "config": cfg, + "avg_answer_relevance": avg_rel, + "cost": float(cost), + "time_s": round(total_time, 2), + "search_s": round(t_search, 2), "gen_s": round(t_gen, 2), "judge_s": round(t_judge, 2), + }) + + logging.info( + f"Trial {tid}: answer_relevance={avg_rel:.4f} cost={cost:.4f} " + f"time={total_time:.1f}s (search={t_search:.1f} gen={t_gen:.1f} judge={t_judge:.1f})" + ) + return float(cost) + + +def param_importance(smac, output_path): + params = list(smac.scenario.configspace.keys()) + configs, costs = [], [] + for key, val in smac.runhistory.items(): + configs.append(dict(smac.runhistory.get_config(key.config_id))) + costs.append(val.cost) + if len(configs) < 3: + return {} + + raw = {} + for p in params: + groups = {} + for c, cost in zip(configs, np.array(costs)): + groups.setdefault(str(c[p]), []).append(cost) + means = [np.mean(g) for g in groups.values()] + raw[p] = float(np.var(means)) if len(means) > 1 else 0.0 + + total = sum(raw.values()) + imp = {p: round(v / total, 4) for p, v in raw.items()} if total > 0 else raw + ranking = sorted(imp, key=imp.get, reverse=True) + result = {"method": "variance_based", "ranking": ranking, "individual": imp} + with open(output_path / "parameter_importance.json", "w") as f: + json.dump(result, f, indent=2) + return result + + +def run(dataset_name, max_questions, n_trials, output_dir): + if dataset_name not in DATASET_CONFIG: + raise ValueError( + f"Unknown dataset: {dataset_name}. " + f"Available: {list(DATASET_CONFIG.keys())}" + ) + + dcfg = DATASET_CONFIG[dataset_name] + out = pathlib.Path(output_dir) / dataset_name + out.mkdir(parents=True, exist_ok=True) + answers_dir = out / "trial_answers" + answers_dir.mkdir(parents=True, exist_ok=True) + + qa = read_dataset_qa(dataset_name, dcfg["subdirectory"], dcfg["split"]) + if max_questions: + qa = qa.select(range(min(max_questions, len(qa)))) + + questions = [] + for i, r in enumerate(qa): + questions.append({ + "question": r.get(dcfg["question_field"], ""), + "ground_truth": r.get(dcfg["answer_field"], ""), + "id": r.get(dcfg["id_field"], str(i)), + }) + logging.info(f"Loaded {len(questions)} questions from {dataset_name}") + + cs = ConfigurationSpace(seed=42) + cs.add([Integer("top_k", (1, 20), default=4), + Integer("chunk_size", (200, 4000), default=2000), + Integer("chunk_overlap", (0, 500), default=200), + Categorical("prompt_template", ["default", "concise", "detailed"], default="default")]) + cs.add(ForbiddenGreaterThanRelation(cs["chunk_overlap"], cs["chunk_size"])) + + scenario = Scenario(configspace=cs, deterministic=True, n_trials=n_trials, + walltime_limit=36000, seed=42, output_directory=out / "smac_output") + + state = {"trial_count": 0, "best_cost": 1.0, "convergence": [], "history": []} + smac = HyperparameterOptimizationFacade( + scenario=scenario, + target_function=lambda config, seed=0: objective( + config, seed, questions, dataset_name, state, answers_dir + ), + ) + + t0 = time.time() + incumbent = smac.optimize() + total_time = time.time() - t0 + + importance = param_importance(smac, out) + + with open(out / "optimization_results.json", "w") as f: + json.dump({ + "best_config": dict(incumbent), + "run_history": state["history"], + "convergence": state["convergence"], + "parameter_importance": importance, + "total_time_s": round(total_time, 2), + "dataset": dataset_name, + "generator": "gemma-3-27b-it", + "judge": "gemma-3-27b-it (self-judge)", + "metrics": ["answer_relevance"], + "metric_source": "Es et al. (2023) RAGAS arXiv:2309.15217v2", + }, f, indent=2, default=str) + + logging.info(f"Done in {total_time:.0f}s") + return incumbent, state["history"], importance + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--dataset", type=str, default="WikiEval", + choices=list(DATASET_CONFIG.keys())) + parser.add_argument("--max_questions", type=int, default=50) + parser.add_argument("--n_trials", type=int, default=10) + parser.add_argument("--output_dir", type=str, default="optimization_results_selfjudge") + args = parser.parse_args() + + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(levelname)s %(message)s", + force=True, + ) + incumbent, history, importance = run( + args.dataset, args.max_questions, args.n_trials, args.output_dir, + ) + + best = min(history, key=lambda x: x["cost"]) + print(f"\nBest config: {dict(incumbent)}") + print(f"Answer Relevance: {best['avg_answer_relevance']:.4f}") + if importance: + print(f"\nParameter importance:") + for i, p in enumerate(importance["ranking"], 1): + print(f" {i}. {p}: {importance['individual'][p]:.4f}") \ No newline at end of file diff --git a/learn2rag/pipeline/config.py b/learn2rag/pipeline/config.py index 1433172..b0677fb 100644 --- a/learn2rag/pipeline/config.py +++ b/learn2rag/pipeline/config.py @@ -1,5 +1,6 @@ import json import os +import logging with open(os.environ.get("PIPELINE_USER_CONFIG", "learn2rag/pipeline/user_config.json"), "r") as file: user_config = json.load(file) @@ -9,3 +10,4 @@ with open(os.environ.get("PIPELINE_OPT_CONFIG", "learn2rag/pipeline/opt_config.json"), "r") as file: opt_config = json.load(file) + logging.info(f"Loaded opt_config:\n{json.dumps(opt_config, indent=4)}") diff --git a/learn2rag/pipeline/debug.py b/learn2rag/pipeline/debug.py new file mode 100755 index 0000000..d1a9d85 --- /dev/null +++ b/learn2rag/pipeline/debug.py @@ -0,0 +1,59 @@ +import asyncio +import logging +import logging.config +import yaml +import json + +from langchain_core.documents.base import Document + +from . import ingestion +from . import search +from . import generate +from .store import delete_collection, delete_documents, get_documents, update_documents + + +if __name__ == "__main__": + + try: + logging.config.dictConfig(yaml.safe_load(open("./learn2rag/pipeline/logging.yaml").read())) + except FileNotFoundError: + logging.basicConfig() + + from .config import user_config, opt_config + + #delete_collection(loader_id="json_test_file", user_config=user_config, opt_config=opt_config) + #results = get_documents(loader_id="json_test_file", user_config=user_config, opt_config=opt_config) + + with open("loaded_documents.json", "r", encoding="utf-8") as f: + raw = json.load(f) + + documents = [ + Document(page_content=d["content"], metadata=d["metadata"]) + for d in raw + ] + #update_documents(loader_id="json_test_file", documents=documents, user_config=user_config, opt_config=opt_config) + # warning: index does not update, use update whereever you can! + ingestion.index(documents, user_config, opt_config) + + if opt_config["query_mode"] == "multi": + # in query_mode 'multi' different querys for each vector in the multi-vector are allowed + multi_query = {"content": "What is USM AI?", "title": "What is USM AI?", "summary": "What is USM AI?", "source_path":"USU/ITSM/"} + results = search.search_multi(multi_query, user_config, opt_config, request_id=None) + points = results.points + # modify the query for generation part + query = " ".join(f"{k}={v}" for k, v in multi_query.items()) + else: + query = "Was sind A, B und C?" + user = "anonymous" + points = asyncio.run(search.search_authorized(query, user, request_id=None)) + + sources = "\n".join(set(point.payload['source'] for point in points)) # type: ignore[index] + + for point in points: + print(f"ID: {point.id}, Path: {point.payload['source']}, Score: {point.score}") # type: ignore[index] + + answer = generate.generate(query, points, opt_config) + + print(query) + print(answer) + print(sources) diff --git a/learn2rag/pipeline/generate.py b/learn2rag/pipeline/generate.py index aa47fbf..5e3a21e 100644 --- a/learn2rag/pipeline/generate.py +++ b/learn2rag/pipeline/generate.py @@ -1,6 +1,7 @@ -from typing import Any, Generator +from typing import Any, Generator, Sequence import logging from langchain.prompts import SystemMessagePromptTemplate, HumanMessagePromptTemplate, ChatPromptTemplate + from qdrant_client.http.models import ScoredPoint from .llm import llm @@ -12,11 +13,15 @@ {content} """ -def generate(query: str, search_results: list[ScoredPoint], opt_config: dict[str, Any]) -> Any: +def generate(query: str, search_results: Sequence[ScoredPoint], opt_config: dict[str, Any]) -> Any: assert llm is not None if hasattr(search_results, "points"): search_results = search_results.points - context = "\n\n".join([context_template.format(source=result.payload['source'], content=result.payload['content']) for result in search_results]) # type: ignore[index] + context = "\n\n".join([ + context_template.format(source=result.payload['source'], content=result.payload['content']) + for result in search_results if result.payload + ]) + system_message = SystemMessagePromptTemplate.from_template(opt_config["prompt"]) user_message = HumanMessagePromptTemplate.from_template("{question}") prompt = ChatPromptTemplate.from_messages([system_message, user_message]) diff --git a/learn2rag/pipeline/ingestion.py b/learn2rag/pipeline/ingestion.py index be2d036..ce688e8 100644 --- a/learn2rag/pipeline/ingestion.py +++ b/learn2rag/pipeline/ingestion.py @@ -6,6 +6,8 @@ import numpy as np import warnings from collections.abc import Iterator +from collections import deque +from time import perf_counter from langchain.text_splitter import RecursiveCharacterTextSplitter from langchain_core.documents import Document @@ -15,6 +17,13 @@ from .embeddings import create_embeddings +def _format_hhmmss(total_seconds: float) -> str: + seconds = max(0, int(total_seconds)) + hours, remainder = divmod(seconds, 3600) + minutes, secs = divmod(remainder, 60) + return f"{hours:02d}:{minutes:02d}:{secs:02d}" + + def get_chunks_metadata(chunks: list[Document], item: str) -> Iterator[str]: missing = 0 for chunk in chunks: @@ -136,8 +145,16 @@ def ingest_batch(docs: list[Document], qdrant: Qdrant, user_config: dict[str, An chunks = text_splitter.split_documents(docs) ingestion_batch_size = opt_config["ingestion_batch_size"] + total_batches = (len(chunks) + ingestion_batch_size - 1) // ingestion_batch_size + eta_window_size = 100 + report_every = 100 + recent_batch_durations: deque[float] = deque(maxlen=eta_window_size) + ingest_start = perf_counter() + logging.info('Creating embeddings and ingesting in batches...') - for batch_start in range(0, len(chunks), ingestion_batch_size): + + for batch_idx, batch_start in enumerate(range(0, len(chunks), ingestion_batch_size), start=1): + batch_started_at = perf_counter() batch_chunks = chunks[batch_start:batch_start + ingestion_batch_size] batch_content = [chunk.page_content for chunk in batch_chunks] # prevent Surrogate Halves errors through invalid UTF-8 characters in the text through replacing @@ -229,7 +246,34 @@ def ingest_batch(docs: list[Document], qdrant: Qdrant, user_config: dict[str, An else: insert(qdrant, collection_name, sample) + batch_duration = perf_counter() - batch_started_at + recent_batch_durations.append(batch_duration) + + if batch_idx % report_every == 0: + elapsed = perf_counter() - ingest_start + avg_batch_duration = sum(recent_batch_durations) / len(recent_batch_durations) + remaining_batches = total_batches - batch_idx + eta_seconds = avg_batch_duration * remaining_batches + progress_percent = (batch_idx / total_batches) * 100 if total_batches > 0 else 100.0 + + logging.info( + "Ingestion progress: %d/%d batches (%.2f%%), elapsed=%s, eta in %s", + batch_idx, + total_batches, + progress_percent, + _format_hhmmss(elapsed), + _format_hhmmss(eta_seconds), + ) + + total_elapsed = perf_counter() - ingest_start + logging.info( + "Ingestion finished: %d/%d batches (100.00%%), total_elapsed=%s", + total_batches, + total_batches, + _format_hhmmss(total_elapsed), + ) + def index(documents: list[Document], user_config: dict[str, Any], opt_config: dict[str, Any]) -> None: """ Ingest a list of documents — entry point for standalone pipeline operation. diff --git a/learn2rag/pipeline/llm.py b/learn2rag/pipeline/llm.py index 0060049..7a031fd 100644 --- a/learn2rag/pipeline/llm.py +++ b/learn2rag/pipeline/llm.py @@ -1,5 +1,6 @@ import logging import os +import httpx from pydantic import SecretStr from langchain_core.language_models.chat_models import BaseChatModel from langchain_core.messages import AIMessage, BaseMessage, SystemMessage @@ -12,6 +13,25 @@ logger = logging.getLogger(__name__) +def _env_float(name: str, default: float) -> float: + value = os.environ.get(name) + if value is None: + return default + try: + return float(value) + except ValueError: + logger.warning("Invalid float for %s=%r. Falling back to %s.", name, value, default) + return default + + +def _ollama_timeout() -> httpx.Timeout: + timeout_s = max(1.0, _env_float("L2R_OLLAMA_TIMEOUT_SECONDS", 90.0)) + connect_s = min(10.0, timeout_s) + write_s = min(30.0, timeout_s) + pool_s = min(10.0, timeout_s) + return httpx.Timeout(timeout=timeout_s, connect=connect_s, read=timeout_s, write=write_s, pool=pool_s) + + class LLMClient(): ID: str '''A key stored in user data, must not be changed''' @@ -63,6 +83,7 @@ def __init__(self, *, url: str, token: str | None, model: str, proxy: str | None client_kwargs={ 'headers': {'Authorization': f'Bearer {token}'} if token else {}, 'proxy': proxy, + 'timeout': _ollama_timeout(), }, ) diff --git a/learn2rag/pipeline/main.py b/learn2rag/pipeline/main.py deleted file mode 100755 index 0e2505b..0000000 --- a/learn2rag/pipeline/main.py +++ /dev/null @@ -1,90 +0,0 @@ -import asyncio -import logging -import logging.config -import yaml -from operator import itemgetter - -from langchain_core.documents.base import Document - -from . import ingestion -from . import search -from . import generate -from .operators import BasicPipeline -from .store import delete_collection, delete_documents, get_documents, update_documents - - -async def main() -> None: - try: - logging.config.dictConfig(yaml.safe_load(open("./learn2rag/pipeline/logging.yaml").read())) - except FileNotFoundError: - logging.basicConfig() - - from .config import user_config, opt_config - - #delete_collection(loader_id="local_docs", user_config=user_config, opt_config=opt_config) - #results = get_documents(loader_id="local_docs", user_config=user_config, opt_config=opt_config) - - documents = [ - Document(page_content=d["content"], metadata=d["metadata"]) - for d in [ - { - "metadata": { - "source": "C:C:\\Users\\foo\\Revised Manuscript_Text categorization approach.docx", - "content_hash": "e18e509d138cf86c22df0b0dfafc5ca5b8f1e266f5e3470de68190f3ebe495b0", - "source_path": "C:\\Users\\foo", - "file_extension": "docx", - "process_date": "2025-07-28", - "process_time": "14:42:02", - "loader_type": "DirectoryLoader", - "loader_id": "local_docs", - "title": "The title of a real document", - "summary": "This document is awesome" - }, - "content": "A brand-new Corpus-based Real-time Text Classification and Tagging Approach for Social Data..." - }, - { - "metadata": { - "source": "C:C:\\Users\\foo\\qdrant.docx", - "content_hash": "7f3b9c1a0d4e6f8b2c5a7d9e1f0b3c6d8a4e2f1c9b7d0a6e5f1c3a8b9d2e4f0", - "source_path": "C:\\Users\\foo", - "file_extension": "docx", - "process_date": "2025-07-28", - "process_time": "14:42:02", - "loader_type": "DirectoryLoader", - "loader_id": "local_docs", - "title": "The title of a real document", - "summary": "This document is awesome" - }, - "content": "Qdrant ist eine Open-Source-Vektordatenbank..." - }, - ] -] - #update_documents(loader_id="local_docs", documents=documents, user_config=user_config, opt_config=opt_config) - ingestion.index(documents, user_config, opt_config) - - if opt_config["query_mode"] == "multi": - # in query_mode 'multi' different querys for each vector in the multi-vector are allowed - multi_query = {"content": "What is USM AI?", "title": "What is USM AI?", "summary": "What is USM AI?", "source_path":"USU/ITSM/"} - results = search.search_multi(multi_query, user_config, opt_config) - points = results.points - # modify the query for generation part - query = " ".join(f"{k}={v}" for k, v in multi_query.items()) - answer = generate.generate(query, points, opt_config) - else: - pipeline = BasicPipeline() - query = "Was sind A, B und C?" - answer, points = itemgetter('answer', 'documents')(await pipeline( - inputs={'question': query, 'user': 'anonymous'}, - )) - - sources = "\n".join(set(point.payload['path'] for point in points)) # type: ignore[index] - - for point in points: - print(f"ID: {point.id}, Path: {point.payload['source']}, Score: {point.score}") # type: ignore[index] - - print(query) - print(answer) - print(sources) - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/learn2rag/pipeline/qdrant.py b/learn2rag/pipeline/qdrant.py index 05bf021..d1f315e 100644 --- a/learn2rag/pipeline/qdrant.py +++ b/learn2rag/pipeline/qdrant.py @@ -17,6 +17,7 @@ class Qdrant: + # FIXME: only create QdrantClient instance when needed, # do not create it when just importing this module client = QdrantClient( @@ -25,6 +26,7 @@ class Qdrant: path=path, ) + def __init__(self, collection_name: str, opt_config: dict[str, Any]) -> None: self.collection_name = collection_name self.vector_size = opt_config["vector_size"][opt_config["embedding_model"]] diff --git a/learn2rag/pipeline/rewrite.py b/learn2rag/pipeline/rewrite.py index 42e405e..e2eb87f 100644 --- a/learn2rag/pipeline/rewrite.py +++ b/learn2rag/pipeline/rewrite.py @@ -1,8 +1,73 @@ from langchain_core.messages import SystemMessage, HumanMessage import ast +import logging +import os +import time from .llm import llm +logger = logging.getLogger(__name__) + + +def _env_int(name: str, default: int) -> int: + value = os.environ.get(name) + if value is None: + return default + try: + return int(value) + except ValueError: + logger.warning("Invalid int for %s=%r. Falling back to %s.", name, value, default) + return default + + +def _env_float(name: str, default: float) -> float: + value = os.environ.get(name) + if value is None: + return default + try: + return float(value) + except ValueError: + logger.warning("Invalid float for %s=%r. Falling back to %s.", name, value, default) + return default + + +def _invoke_llm(messages: list[SystemMessage | HumanMessage], *, purpose: str) -> str: + if llm is None: + return '' + + max_attempts = max(1, _env_int("L2R_OLLAMA_MAX_RETRIES", 2)) + retry_sleep_s = max(0.0, _env_float("L2R_OLLAMA_RETRY_BACKOFF_SECONDS", 2.0)) + + for attempt in range(1, max_attempts + 1): + try: + t0 = time.time() + logger.info("llm_invoke_start purpose=%s attempt=%d/%d", purpose, attempt, max_attempts) + response = llm.invoke(messages, stream=False) + duration_s = time.time() - t0 + logger.info( + "llm_invoke_done purpose=%s attempt=%d/%d duration_s=%.2f", + purpose, + attempt, + max_attempts, + duration_s, + ) + content = response.content + return content.strip() if isinstance(content, str) else '' + except Exception as exc: + logger.warning( + "llm_invoke_failed purpose=%s attempt=%d/%d error=%s", + purpose, + attempt, + max_attempts, + exc, + ) + if attempt < max_attempts and retry_sleep_s > 0: + time.sleep(retry_sleep_s * attempt) + + logger.error("llm_invoke_give_up purpose=%s attempts=%d", purpose, max_attempts) + return '' + + # future todo: add history handling / add state handling for loops in pipeline try: @@ -25,15 +90,11 @@ def rewrite_query(user_query: str) -> str: - Return only the rewritten query text. """ - response = llm.invoke([ + content = _invoke_llm([ SystemMessage(content=system_message_rewrite_query), HumanMessage(content=user_query), - ]) - - content = response.content - if isinstance(content, str): - return content.strip() - return '' + ], purpose="rewrite_query") + return content def generate_subqueries(user_query: str, n: int=3) -> list[str]: @@ -53,21 +114,19 @@ def generate_subqueries(user_query: str, n: int=3) -> list[str]: {synonym_list} """ - response = llm.invoke([ + content = _invoke_llm([ SystemMessage(content=system_message_generate_subqueries), HumanMessage(content=user_query), - ]) - - content = response.content - if not isinstance(content, str): + ], purpose="generate_subqueries") + if not content: return [] try: result = ast.literal_eval(content.strip()) if isinstance(result, list): return [str(x).strip() for x in result if str(x).strip()] - except Exception: - pass + except Exception as exc: + logger.warning("generate_subqueries_parse_failed query=%r error=%s content=%r", user_query, exc, content[:500]) return [] @@ -91,20 +150,18 @@ def generate_keywords(user_query: str, n: int=3) -> list[str]: {synonym_list} """ - response = llm.invoke([ + content = _invoke_llm([ SystemMessage(content=system_message_generate_keywords), HumanMessage(content=user_query), - ]) - - content = response.content - if not isinstance(content, str): + ], purpose="generate_keywords") + if not content: return [] try: result = ast.literal_eval(content.strip()) if isinstance(result, list): return [str(x).strip() for x in result if str(x).strip()] - except Exception: - pass + except Exception as exc: + logger.warning("generate_keywords_parse_failed query=%r error=%s content=%r", user_query, exc, content[:500]) return [] diff --git a/learn2rag/pipeline/search.py b/learn2rag/pipeline/search.py index a4b43ce..71b60e7 100644 --- a/learn2rag/pipeline/search.py +++ b/learn2rag/pipeline/search.py @@ -28,7 +28,6 @@ def _get_flag_reranker(model_name: str, use_fp16: bool) -> FlagReranker: def _get_cross_encoder(model_name: str) -> CrossEncoder: return cast(CrossEncoder, CrossEncoder(model_name)) - def _sort_and_deduplicate(points: list[ScoredPoint]) -> list[ScoredPoint]: best_by_id: dict[str, ScoredPoint] = {} fallback_points: list[ScoredPoint] = [] @@ -118,6 +117,7 @@ def _rerank_points_with_colbert( *, top_k: int, opt_config: dict[str, Any], + user_config: dict[str, Any], ) -> list[ScoredPoint]: collection_name = user_config["collection_name"] qdrant = Qdrant(collection_name=collection_name, opt_config=opt_config) @@ -137,6 +137,7 @@ def _rerank_points_with_colbert( query=colbert_query, # type: ignore[arg-type] using="colbert", limit=top_k, + timeout=120 ) return list(results.points) @@ -195,8 +196,28 @@ def _collect_query_points( extra={'activity': '_collect_query_points', 'request_id': request_id}, ) - for sq in subqueries: - sq_results = search(sq, user_config, opt_config_subqueries) + + for idx, sq in enumerate(subqueries, start=1): + profilingLogger.info( + "subquery_search_start query=%r subquery_index=%d/%d subquery=%r top_k=%s", + query, + idx, + len(subqueries), + sq, + opt_config_subqueries["top_k"], + extra={'activity': '_collect_query_points', 'request_id': request_id}, + ) + sq_results = search(sq, user_config, opt_config_subqueries,request_id=request_id) + profilingLogger.info( + "subquery_search_done query=%r subquery_index=%d/%d subquery=%r points=%d", + query, + idx, + len(subqueries), + sq, + len(sq_results.points), + extra={'activity': '_collect_query_points', 'request_id': request_id}, + ) + points_all.extend(sq_results.points) if rewrite_mode in ["keywords", "subqueries_keywords"]: @@ -215,8 +236,28 @@ def _collect_query_points( extra={'activity': '_collect_query_points', 'request_id': request_id}, ) - for kw in keywords: - kw_results = search(kw, user_config, opt_config_keywords) + + for idx, kw in enumerate(keywords, start=1): + profilingLogger.info( + "keyword_search_start query=%r keyword_index=%d/%d keyword=%r top_k=%s", + query, + idx, + len(keywords), + kw, + opt_config_keywords["top_k"], + extra={'activity': '_collect_query_points', 'request_id': request_id}, + ) + kw_results = search(kw, user_config, opt_config_keywords, request_id=request_id) + profilingLogger.info( + "keyword_search_done query=%r keyword_index=%d/%d keyword=%r points=%d", + query, + idx, + len(keywords), + kw, + len(kw_results.points), + extra={'activity': '_collect_query_points', 'request_id': request_id}, + ) + points_all.extend(kw_results.points) points = _sort_and_deduplicate(points_all) @@ -249,21 +290,23 @@ def _collect_query_points( query, points, top_k=opt_config["top_k_reranker"], - opt_config=opt_config + opt_config=opt_config, + user_config=user_config, ) - + else: + points = points[:opt_config["top_k"]] return points # similarity search -def search(query: str, user_config: dict[str, Any], opt_config: dict[str, Any]) -> QueryResponse: - profilingLogger.info('start', extra={'activity': 'search'}) +def search(query: str, user_config: dict[str, Any], opt_config: dict[str, Any], *, request_id: str | None = None) -> QueryResponse: + profilingLogger.info('start', extra={'activity': 'search', 'request_id': request_id}) profilingLogger.info( "search_called query=%r search_mode=%s collection_name=%s", query, opt_config.get("search_mode"), user_config.get("collection_name"), - extra={'activity': '_collect_query_points'}, + extra={'activity': '_collect_query_points', 'request_id': request_id}, ) collection_name = user_config["collection_name"] @@ -317,6 +360,7 @@ def search(query: str, user_config: dict[str, Any], opt_config: dict[str, Any]) query=query_embedding, # type: ignore[arg-type, unused-ignore] using="dense", limit=opt_config["top_k"], + timeout=120 ) elif opt_config["search_mode"] == "sparse": indices = [int(k) for k in query_embedding["lexical_weights"].keys()] # type: ignore[union-attr] @@ -326,6 +370,7 @@ def search(query: str, user_config: dict[str, Any], opt_config: dict[str, Any]) query=models.SparseVector(indices=indices, values=values), using="sparse", limit=opt_config["top_k"], + timeout=120 ) elif opt_config["search_mode"] == "dense_sparse": indices = [int(k) for k in query_embedding["lexical_weights"].keys()] # type: ignore[union-attr] @@ -346,6 +391,7 @@ def search(query: str, user_config: dict[str, Any], opt_config: dict[str, Any]) ], query=models.FusionQuery(fusion=fusion_mode), limit=opt_config["top_k"], + timeout=120 ) elif opt_config["search_mode"] == "dense_sparse_colbert": @@ -372,6 +418,7 @@ def search(query: str, user_config: dict[str, Any], opt_config: dict[str, Any]) ], query=models.FusionQuery(fusion=fusion_mode), limit=opt_config["top_k"], + timeout=120 ) elif opt_config["search_mode"] == "multi_search": @@ -380,6 +427,7 @@ def search(query: str, user_config: dict[str, Any], opt_config: dict[str, Any]) query=query_embedding, # type: ignore[arg-type, unused-ignore] using="multi", limit=opt_config["top_k"], + timeout=120 ) return results @@ -413,11 +461,15 @@ def search_multi(multi_query: dict[str, str], user_config: dict[str, Any], opt_c query=query_embedding, # type: ignore[arg-type, unused-ignore] using="multi", limit=opt_config["top_k"], + timeout=120 ) return results -async def search_authorized(question: str, user: str) -> List[ScoredPoint]: - points = _collect_query_points(question, user_config, opt_config) + + +async def search_authorized(question: str, user: str, *, request_id: str | None = None, user_config: dict[str, Any] = user_config, opt_config: dict[str, Any] = opt_config) -> List[ScoredPoint]: + points = _collect_query_points(question, user_config, opt_config, request_id=request_id) + query_response = QueryResponse(points=points) authorized_points = await filter_authorized(user, query_response) # keep deterministic order after auth filter diff --git a/learn2rag/tests/data/rabbit_eval.csv b/learn2rag/tests/data/rabbit_eval.csv new file mode 100644 index 0000000..df81046 --- /dev/null +++ b/learn2rag/tests/data/rabbit_eval.csv @@ -0,0 +1,7 @@ +answer;question +"Answer: Rabbits belong to the family Leporidae, which also includes hares, and the order Lagomorpha, which includes pikas.";"Question: What family and order do rabbits belong to, and what other animals are included in these classifications?" +"Answer: The most widespread rabbit genera are Oryctolagus and Sylvilagus. The European rabbit, which belongs to the genus Oryctolagus, has been introduced on every continent except Antarctica.";"Question: What are the most widespread rabbit genera, and on which continents can the European rabbit be found?" +"Answer: Rabbits do not constitute a clade because they are a paraphyletic grouping, as hares are nested within the Leporidae clade but not described as rabbits. They differ from rodents by having a number of traits rodents lack, including two extra incisors.";"Question: Why do rabbits not constitute a clade, and how do they physically differ from rodents despite once being classified as such?" +"Answer: A rabbit's hind legs are longer than its fore legs, allowing for quick hopping to escape predators and providing powerful kicks if captured. Rabbits are typically nocturnal and often sleep with their eyes open.";"Question: How does the bone structure of a rabbit's legs aid in its survival against predators, and what are their typical sleep habits?" +"Answer: A rabbit's ears are essential for thermoregulation, containing a high density of blood vessels. They also have a high surface area to help detect potential predators.";"Question: What functions do a rabbit's ears serve regarding survival and physiology?" +"Answer: Humans have used rabbits as livestock since at least the first century BC in ancient Rome, raising them for meat, fur, and wool. The practice of raising and breeding them is known as cuniculture.";"Question: When did humans begin using rabbits as livestock, for what purposes, and what is the practice of breeding them called?" \ No newline at end of file diff --git a/learn2rag/tests/test_learn2rag.py b/learn2rag/tests/test_learn2rag.py index dab7d1b..acda46a 100644 --- a/learn2rag/tests/test_learn2rag.py +++ b/learn2rag/tests/test_learn2rag.py @@ -14,6 +14,7 @@ logger = logging.getLogger(__name__) template_dir = Path(__file__).resolve().parent.parent / 'ui' / 'templates' / 'compose' / 'pipelines' +# optimization_dir = Path(__file__).resolve().parent.parent / 'optimization' data_dir = Path(__file__).resolve().parent / 'data' @@ -112,3 +113,110 @@ def check_rag() -> None: except APIConnectionError: assert False waitUntil(check_rag, timeout=1 * 60 * 1000) + + def test_optimization(self) -> None: + template_context = { + 'is_windows': is_windows(), + 'learn2rag_path': Path('.').absolute(), + 'storage_path': self.storage_path, + 'ports': { + 'pipeline': self.rag_port, + }, + 'qdrant_api_key': '', + 'language_model': {'api': 'ChatFake'}, + 'pipeline': { + 'qdrant_path': self.storage_path / 'qdrant_persistence', + }, + 'import_config': { + 'loaders': [ + { + 'loader_id': 'local_test', + 'loader_type': 'DirectoryLoader', + 'recursive': 'True', + 'path': str(data_dir), + }, + ], + }, + } + + project = Project.create(template_dir / 'import.yml', self.project_name, template=True, + template_context=template_context) + assert project is not None, 'project should not be None' + project.start() + assert project.running + + def check_import() -> None: + project = Project.get(self.project_name) + assert project is not None + assert not project.running + + waitUntil(check_import, timeout=1 * 60 * 1000) + + project.remove() + + project = Project.create(template_dir / 'pipeline.yml', self.project_name, template=True, + template_context=template_context) + + assert project is not None + project.start() + + def check_pipeline() -> None: + try: + self.openai_client.models.list() + except APIConnectionError: + assert False + + waitUntil(check_pipeline, timeout=1 * 60 * 1000) + + # Optimization + from learn2rag.optimization.baseline_optimization import run + import os + import json + dataset_name = "test_rabbit" + opt_out_dir = self.storage_path / "opt_output" + results_file = opt_out_dir / dataset_name / "optimization_results.json" + + mock_user_config = { + "collection_name": dataset_name, + "qdrant_path": str(self.storage_path / 'qdrant_persistence') + } + os.environ["PIPELINE_USER_CONFIG"] = json.dumps(mock_user_config) + + initial_mtime = results_file.stat().st_mtime if results_file.exists() else 0.0 + + + mock_registry = { + "datasets": { + dataset_name: { + "subdirectory": "", "split": "train", + "fields": {"q": "question", "a": "answer", "id": "id"}, + "path": str(data_dir / "rabbit_eval.csv") + } + }, + "prompts": { + "default": "Answer using ONLY the provided information: {context}", + "concise": "Be concise. Information: {context}" + } + } + + best_cfg, history, importance = run( + dataset_name=dataset_name , + max_questions=2, + n_trials=2, + output_dir=opt_out_dir, + registry_path=mock_registry + ) + + + assert best_cfg is not None, "Optimization should return a valid configuration" + assert len(history) == 2, "History length should match n_trials" + + + assert results_file.exists(), "Optimization output JSON was not created" + current_mtime = results_file.stat().st_mtime + assert current_mtime > initial_mtime, "The optimization results file was not updated during the run!" + + with open(results_file, 'r') as f: + results_data = json.load(f) + assert "best_config" in results_data + assert "top_k" in results_data["best_config"], "Optimization failed to output expected parameters" \ No newline at end of file diff --git a/learn2rag/ui/__init__.py b/learn2rag/ui/__init__.py index b235a00..c69530f 100644 --- a/learn2rag/ui/__init__.py +++ b/learn2rag/ui/__init__.py @@ -12,6 +12,7 @@ import time from typing import Any import urllib +from itertools import islice from babel import negotiate_locale from flask import Flask, flash, redirect as flask_redirect, render_template, request, make_response, url_for @@ -23,6 +24,7 @@ import werkzeug.wrappers from learn2rag.compose import Project +from learn2rag.evaluation.tools import read_dataset_qa import learn2rag.data import learn2rag.pipeline.llm from ..utils import ( @@ -32,7 +34,7 @@ save_data_path, ) -from datetime import datetime # <-- ADD THIS +from datetime import datetime logging.getLogger().addHandler(flask.logging.default_handler) @@ -406,7 +408,6 @@ def pipeline_create() -> 'str | werkzeug.wrappers.response.Response': def start_pipeline(name: str, pipeline: dict[str, Any], template_name: str) -> None: has_ssl = bool(app.config.get("TLS")) - url = urllib.parse.urlparse(request.base_url) assert url.scheme @@ -480,6 +481,41 @@ def start_pipeline(name: str, pipeline: dict[str, Any], template_name: str) -> N # TODO "load" the corresponding Ollama model + @app.get('/pipelines/') + def pipeline_details(name: str) -> 'str | werkzeug.wrappers.response.Response': + pipeline = learn2rag.data.get_entry(app.instance_path, 'pipelines', name) + if pipeline is None: + flash(pgettext('flash', 'The requested pipeline is not found'), 'error') + return redirect(url_for('pipelines_list')) + storage_path = Path(pipeline['storage_path']) + try: + training_dataset = read_dataset_qa(storage_path / 'training.csv', 'train') + except FileNotFoundError: + training_dataset = None + return render_template( + 'pipelines_details_page.html', + name=name, + pipeline=pipeline, + training_dataset=training_dataset, + projects=Project.get_all(), + ) + + @app.post('/pipelines//training') + def pipeline_details_training_data(name: str) -> 'str | werkzeug.wrappers.response.Response': + pipeline = learn2rag.data.get_entry(app.instance_path, 'pipelines', name) + if pipeline is None: + flash(pgettext('flash', 'The requested pipeline is not found'), 'error') + return redirect(url_for('pipelines_list')) + try: + storage_path = Path(pipeline['storage_path']) + storage_path.mkdir(parents=True, exist_ok=True) + training_file = request.files['trainingFile'] + training_file.save(storage_path / 'training.csv') + except Exception as e: + app.logger.exception(e) + flash(pgettext('flash', 'Could not save the file'), 'error') + return redirect(url_for('pipeline_details', name=name)) + @app.post('/pipelines/') def pipeline_action(name: str) -> 'str | werkzeug.wrappers.response.Response': pipeline = learn2rag.data.get_entry(app.instance_path, 'pipelines', name) diff --git a/learn2rag/ui/templates/compose/pipelines/continuous.yml b/learn2rag/ui/templates/compose/pipelines/continuous.yml index 81981a1..6d6f5a7 100644 --- a/learn2rag/ui/templates/compose/pipelines/continuous.yml +++ b/learn2rag/ui/templates/compose/pipelines/continuous.yml @@ -29,6 +29,10 @@ files: "llm": "{{language_model.model}}" } + - path: '{{storage_path}}/opt_config.json' + src: '{{learn2rag_path}}/learn2rag/pipeline/opt_config.json' + force: false + - path: '{{storage_path}}/logging_config.yml' content: | version: 1 @@ -193,7 +197,7 @@ services: QDRANT__SERVICE__API_KEY: '{{qdrant_api_key}}' PIPELINE_USER_CONFIG: '{{storage_path}}/basic_user_config.json' IMPORTER_CONFIG: '{{storage_path}}/importer_config.json' - PIPELINE_OPT_CONFIG: '{{learn2rag_path}}/learn2rag/pipeline/opt_config.json' + PIPELINE_OPT_CONFIG: '{{storage_path}}/opt_config.json' LANGCHAIN_API_KEY: '1' LANGCHAIN_TRACING_V2: 'false' LLM_API_TYPE: '{{language_model.api}}' @@ -215,7 +219,7 @@ services: QDRANT__SERVICE__API_KEY: '{{qdrant_api_key}}' PIPELINE_USER_CONFIG: '{{storage_path}}/basic_user_config.json' IMPORTER_CONFIG: '{{storage_path}}/importer_config.json' - PIPELINE_OPT_CONFIG: '{{learn2rag_path}}/learn2rag/pipeline/opt_config.json' + PIPELINE_OPT_CONFIG: '{{storage_path}}/opt_config.json' LANGCHAIN_API_KEY: '1' LANGCHAIN_TRACING_V2: 'false' LLM_API_TYPE: '{{language_model.api}}' diff --git a/learn2rag/ui/templates/compose/pipelines/import.yml b/learn2rag/ui/templates/compose/pipelines/import.yml index 6393279..b81825d 100644 --- a/learn2rag/ui/templates/compose/pipelines/import.yml +++ b/learn2rag/ui/templates/compose/pipelines/import.yml @@ -15,6 +15,10 @@ files: host: '127.0.0.1' telemetry_disabled: true + - path: '{{storage_path}}/opt_config.json' + src: '{{learn2rag_path}}/learn2rag/pipeline/opt_config.json' + force: false + - path: '{{storage_path}}/importer_config.json' content: '{{import_config | tojson}}' @@ -127,7 +131,7 @@ services: QDRANT__SERVICE__API_KEY: '{{qdrant_api_key}}' PIPELINE_USER_CONFIG: '{{storage_path}}/basic_user_config.json' IMPORTER_CONFIG: '{{storage_path}}/importer_config.json' - PIPELINE_OPT_CONFIG: '{{learn2rag_path}}/learn2rag/pipeline/opt_config.json' + PIPELINE_OPT_CONFIG: '{{storage_path}}/opt_config.json' LANGCHAIN_API_KEY: '1' LANGCHAIN_TRACING_V2: 'false' LLM_API_TYPE: '{{language_model.api}}' diff --git a/learn2rag/ui/templates/compose/pipelines/optimization.yml b/learn2rag/ui/templates/compose/pipelines/optimization.yml new file mode 100644 index 0000000..7d36f1d --- /dev/null +++ b/learn2rag/ui/templates/compose/pipelines/optimization.yml @@ -0,0 +1,145 @@ +name: optimization +label: Optimize +ports: + # TODO: labels in the interface currently assume a specific port order + - ui + - qdrant_http +files: + - path: '{{storage_path}}/importer_config.json' + content: '{{import_config | tojson}}' + + - path: '{{storage_path}}/basic_user_config.json' + content: | + { + "collection_name": "learn2rag", + "imported_documents_file_path": "loaded_documents.json", + "llm": "{{language_model.model}}" + } + - path: '{{storage_path}}/opt_config.json' + src: '{{learn2rag_path}}/learn2rag/pipeline/opt_config.json' + force: false + - path: '{{storage_path}}/qdrant_config.yml' + content: | + log_level: ERROR + service: + api_key: '{{qdrant_api_key}}' + grpc_port: null + http_port: {{ports.qdrant_http}} + host: '127.0.0.1' + telemetry_disabled: true + - path: '{{storage_path}}/logging_config.yml' + content: | + version: 1 + formatters: + simple: + format: "%(asctime)s %(levelname)-8s %(name)s %(message)s" + colored: + class: colorlog.ColoredFormatter + format: "%(log_color)s%(asctime)s %(levelname)-8s %(name)s %(message)s" + profiling: + format: "%(created)f %(request_id)s %(activity)s %(message)s" + defaults: + request_id: null + handlers: + display: + class: logging.StreamHandler + level: INFO + formatter: colored + stream: ext://sys.stderr + profiling_display: + class: logging.StreamHandler + formatter: profiling + stream: ext://sys.stderr + profiling_file: + class: logging.FileHandler + formatter: profiling + filename: '{{storage_path}}/logs/profiling.log' + encoding: utf-8 + errors_file: + class: logging.FileHandler + level: ERROR + formatter: simple + filename: '{{storage_path}}/logs/error.log' + encoding: utf-8 + {% if debug_logging %} + debug_file: + class: logging.FileHandler + level: DEBUG + formatter: simple + filename: '{{storage_path}}/logs/debug.log' + encoding: utf-8 + {% endif %} + loggers: + profiling: + handlers: + - profiling_display + - profiling_file + propagate: no + root: + level: DEBUG + handlers: + - display + - errors_file + {% if debug_logging %} + - debug_file + {% endif %} + + - path: '{{storage_path}}/logs/.keep' + content: '' + + - path: '{{storage_path}}/optimization_registry.json' + content: | + { + "datasets": { + "training": { + "subdirectory": "", "split": "train", + "fields": {"q": "question", "a": "answer", "id": "id"}, + "path": "{{storage_path}}/training.csv" + } + }, + "prompts": { + "default": + "# Role and Objective\nYou will act as a smart AI chatbot that answers questions only by using the content from the provided information list.\n\n # Instructions\n- Respond in the language of the question.\n - Answer clear and concise.\n- Only use the provided information.\n - NEVER use your general knowledge.\n\n # Information:\n{context}" + , + "concise": + "Answer the question using ONLY the provided information. Be concise and direct. If the information does not contain the answer, say so.\n\n Information:\n{context}" + , + "detailed": + "You are a knowledgeable assistant. Using ONLY the provided information below, answer the question thoroughly. Cite your sources. If the information is insufficient, state that clearly.\n\n Information:\n{context}" + } + } + +services: + qdrant: + working_dir: '{{storage_path}}' + command: + - '{{learn2rag_path}}/services/qdrant/qdrant{% if is_windows %}.exe{% endif %}' + - '--config-path' + - '{{storage_path}}/qdrant_config.yml' + main: + working_dir: '{{storage_path}}' + command: + - '{{learn2rag_path}}/configurator{% if is_windows %}.exe{% endif %}' + - 'learn2rag.optimization' + - '--logging-config' + - '{{storage_path}}/logging_config.yml' + - '--registry-path' + - '{{storage_path}}/optimization_registry.json' + - '--dataset' + - 'training' + - '--strategy' + - 'retrieval' + environment: + LEARN2RAG_PATH: '{{learn2rag_path}}' + QDRANT_LOCATION: '{{ pipeline.qdrant_location or "http://localhost:" ~ ports.qdrant_http }}' + QDRANT_PATH: '{{ pipeline.qdrant_path }}' + QDRANT__SERVICE__API_KEY: '{{qdrant_api_key}}' + PIPELINE_USER_CONFIG: '{{storage_path}}/basic_user_config.json' + IMPORTER_CONFIG: '{{storage_path}}/importer_config.json' + PIPELINE_OPT_CONFIG: '{{storage_path}}/opt_config.json' + LANGCHAIN_API_KEY: '1' + LANGCHAIN_TRACING_V2: 'false' + LLM_API_TYPE: '{{language_model.api}}' + LLM_API_URL: '{{language_model.url}}' + LLM_API_TOKEN: '{{language_model.token}}' + LLM_API_MODEL: '{{language_model.model}}' diff --git a/learn2rag/ui/templates/compose/pipelines/pipeline.yml b/learn2rag/ui/templates/compose/pipelines/pipeline.yml index b113317..4122257 100644 --- a/learn2rag/ui/templates/compose/pipelines/pipeline.yml +++ b/learn2rag/ui/templates/compose/pipelines/pipeline.yml @@ -24,6 +24,9 @@ files: "imported_documents_file_path": "loaded_documents.json", "llm": "{{language_model.model}}" } + - path: '{{storage_path}}/opt_config.json' + src: '{{learn2rag_path}}/learn2rag/pipeline/opt_config.json' + force: false # importer_config.json should already be present @@ -186,7 +189,7 @@ services: QDRANT__SERVICE__API_KEY: '{{qdrant_api_key}}' PIPELINE_USER_CONFIG: '{{storage_path}}/basic_user_config.json' IMPORTER_CONFIG: '{{storage_path}}/importer_config.json' - PIPELINE_OPT_CONFIG: '{{learn2rag_path}}/learn2rag/pipeline/opt_config.json' + PIPELINE_OPT_CONFIG: '{{storage_path}}/opt_config.json' LANGCHAIN_API_KEY: '1' LANGCHAIN_TRACING_V2: 'false' LLM_API_TYPE: '{{language_model.api}}' diff --git a/learn2rag/ui/templates/pipelines_details_page.html b/learn2rag/ui/templates/pipelines_details_page.html new file mode 100644 index 0000000..fada205 --- /dev/null +++ b/learn2rag/ui/templates/pipelines_details_page.html @@ -0,0 +1,54 @@ +{% extends 'base.html' %} + +{% block header %} +

{% block title %}{{pipeline.label}}{% endblock %}

+{% endblock %} + +{% block content %} +

{{pgettext('header', 'Details')}}

+{{pipeline.label}} +
+

{{pgettext('header', 'Optimization')}}

+{% if training_dataset %} +
+ + + + + + + + + + {% for training_item in training_dataset.select(range(3)) %} + + + + + {% endfor %} + + + + + + + + +
{{pgettext('header', 'Training data')}}
{{pgettext('header', 'Question')}}{{pgettext('header', 'Answer')}}
{{training_item['question']}}{{training_item['answer']}}
{{gettext('Total rows: %(amount)s', amount=training_dataset.num_rows)}}
+
+{% endif %} +{% if training_dataset and not (projects[name] and projects[name].running) %} +
+ +
+{% endif %} +
+ +
+ +
+
+ +
+
+{% endblock %} diff --git a/learn2rag/ui/templates/pipelines_list.html b/learn2rag/ui/templates/pipelines_list.html index 8b5100e..f6fa0fb 100644 --- a/learn2rag/ui/templates/pipelines_list.html +++ b/learn2rag/ui/templates/pipelines_list.html @@ -83,9 +83,15 @@ {{pgettext('button', 'More')}}