import os import torch import numpy as np from typing import List, Dict, Set from src.until.page_similarity_graph import construct_page_similarity_graph class DocumentHypergraph: def __init__(self): self.page_similarity_graph = {} self.query_specific_hypergraph = {} self.graph = {} def construct_page_similarity_graph( self, page_embeds, threshold: float = 0.7, k_value: int = 5, similarity_measure: str = "cosine" ) -> Dict[int, List[int]]: self.page_similarity_graph = construct_page_similarity_graph( page_embeds=page_embeds, threshold=threshold, k_value=k_value, similarity_measure=similarity_measure ) def _get_specific_page( self, basis_pages: List[int], target_pages: List[int] ) -> Set[int]: set_basis, set_target = set(basis_pages), set(target_pages) new_pages = set_target - set_basis improved_pages = { p for p in set_basis & set_target if target_pages.index(p) < basis_pages.index(p) } return new_pages | improved_pages def construct_query_specific_hypergraph( self, bandit, vlm, dataset, top_page_indices: List[List[int]], top_page_scores: List[List[float]], queries: List[str], col_score_dict: Dict[int, float], sample: Dict, best_first_hit_score: int = 20 ): document_hypergraph = {} extra_pages = set() default_idx = len(top_page_indices) - 1 bf_query_idx, bf_vlm_score, bf_score, bf_page = default_idx, -999, -999, 0 def _load_node( page_num: int, query_idx: int, use_vlm: bool = True, first_hit: int = 1, multi_hit: int = 1 ): nonlocal bf_query_idx, bf_vlm_score, bf_score, bf_page if page_num in document_hypergraph: node = document_hypergraph[page_num] node["hits"] += multi_hit cb_score = bandit.sample(node.get("query_idx", set())) node["score"] = bandit._compute_score( col_score=node["col_score"], vlm_score=node["vlm_score"], cb_score=0.5, hits=node["hits"], ) if query_idx != default_idx and query_idx not in node["query_idx"]: node["query_idx"].add(query_idx) node["query"] += f",{queries[query_idx]}" else: vlm_score = -1 if use_vlm: vlm_score = bandit.query_vlm_relevance( vlm=vlm, dataset=dataset, sample=sample, page=page_num, priori=queries[query_idx], ) bandit.update({query_idx}, vlm_score) specific_pages = self._get_specific_page( top_page_indices[default_idx], top_page_indices[query_idx] ) if query_idx != default_idx else set(top_page_indices[default_idx]) neighbors = specific_pages | set(self.page_similarity_graph.get(page_num, [])) extra_pages.update(neighbors - specific_pages) cb_score = bandit.sample({query_idx}) document_hypergraph[page_num] = { "query_idx": {query_idx}, "query": queries[query_idx], "col_score": col_score_dict[page_num], "vlm_score": vlm_score, "hits": first_hit, "cb_score": cb_score, "score": bandit._compute_score(col_score_dict[page_num], vlm_score, 0.5, first_hit), "neighbor": neighbors, } if use_vlm: _vlm = document_hypergraph[page_num]['vlm_score'] _col = document_hypergraph[page_num]['col_score'] print(f"{queries[query_idx]} [page:{page_num}]: \n - vlm: {_vlm}\n - col: {_col}") if (vlm_score > bf_vlm_score or (vlm_score == bf_vlm_score and document_hypergraph[page_num]["score"] > bf_score + 0.0001)): bf_query_idx, bf_vlm_score, bf_score, bf_page = query_idx, vlm_score, document_hypergraph[page_num]["score"], page_num for query_idx in reversed(range(len(top_page_indices))): _load_node( page_num=top_page_indices[query_idx][0], query_idx=query_idx, use_vlm=True, first_hit=best_first_hit_score, ) if bf_page in document_hypergraph: print(f" [DEBUG] best_first_queries: {queries[bf_query_idx]}") sample["best_first_queries"] = queries[bf_query_idx] document_hypergraph[bf_page]["hits"] += best_first_hit_score cb_score = bandit.sample(document_hypergraph[bf_page]["query_idx"]) document_hypergraph[bf_page]["score"] = bandit._compute_score( col_score=document_hypergraph[bf_page]["col_score"], vlm_score=document_hypergraph[bf_page]["vlm_score"], cb_score=cb_score, hits=document_hypergraph[bf_page]["hits"], ) for page_num in top_page_indices[bf_query_idx][:10]: _load_node( page_num=page_num, query_idx=bf_query_idx, use_vlm=False, first_hit=best_first_hit_score ) for query_idx in range(len(top_page_indices)): for page_num in top_page_indices[query_idx]: _load_node( page_num=page_num, query_idx=query_idx, use_vlm=False, first_hit=1 ) for page_num in extra_pages: if page_num not in document_hypergraph: document_hypergraph[page_num] = { "query_idx": set(), "query": queries[default_idx], "col_score": col_score_dict[page_num], "hits": 1, "cb_score": bandit.sample({default_idx}), "neighbor": set(), } self.query_specific_hypergraph = document_hypergraph def _evaluate_rag_one_sample(self, gt, pred, top_k=[1, 3, 5]): metrics = 0 len_gt = len(gt) for k in top_k: cur_pred = pred[:k] intersect = len(set(cur_pred) & set(gt)) metrics += intersect / len_gt * 100.0 metrics += intersect / len(cur_pred) * 100.0 return int(metrics) def _debug(self, dataset, sample, queries, scores): if "evidence_pages" not in sample: print("[Warning]: No ground_truth.") return ground_truth = sample["evidence_pages"] if isinstance(ground_truth, str): try: ground_truth = [int(page.strip()) for page in ground_truth.strip('[]').split(',')] if ground_truth != "[]" else [] except (ValueError, AttributeError): print("Error parsing evidence_pages string") ground_truth = [] else: ground_truth = list(ground_truth) if ground_truth != [] and len(ground_truth) > 0: if dataset.dataset_name == "MMLongBench": ground_truth = [g - 1 if g > 0 else g for g in ground_truth] else: ground_truth = [g if g > 0 else g for g in ground_truth] query_top = torch.topk(torch.tensor(np.array(scores)), min(10, len(scores[0])), dim=-1) query_top_indices = query_top.indices.tolist() query_top_scores = query_top.values.tolist() metrics_list = [] for i, top_indices in enumerate(query_top_indices): metrics = self._evaluate_rag_one_sample(ground_truth, top_indices) metrics_list.append((metrics, queries[i], query_top_indices[i], query_top_scores[i])) metrics_list.sort(key=lambda x: x[0], reverse=True) sample["sorted_queries"] = {query: metric for metric, query, _, _ in metrics_list} print("\n", "#" * 30) print("[DEBUG] Ground Truth: ", ground_truth) result = "".join([f" [{query}, {metric}]\n" for metric, query, _, _ in metrics_list]) print("[DEBUG] True top Indices and Scores: \n" + result) print("#" * 30, "\n") def clean_up_page_similarity_graph(self): self.page_similarity_graph = {} def clean_up_query_specific_hypergraph(self): self.query_specific_hypergraph = {}