Upload run_airbench.py

run_airbench.py

@@ -0,0 +1,105 @@
+from typing import Any, Callable, Dict, List, Optional
+
+from air_benchmark import AIRBench, Retriever
+from llama_index.core import VectorStoreIndex
+from llama_index.core.node_parser import SentenceSplitter
+from llama_index.embeddings.openai import OpenAIEmbedding
+from llama_index.llms.openai import OpenAI
+from llama_index.retrievers.bm25 import BM25Retriever
+from llama_index.core.retrievers import QueryFusionRetriever
+from llama_index.core.schema import Document, NodeWithScore
+
+
+def create_retriever_fn(documents: List[Document], top_k: int) -> Callable[[str], List[NodeWithScore]]:
+    # IMPORTANT: if you don't use a llama-index node parser/splitter, you need to ensure
+    # that node.ref_doc_id points to the correct parent document id.
+    # This is used to line up the corpus document id for evaluation
+    nodes = SentenceSplitter(chunk_size=1024, chunk_overlap=128)(documents)
+
+    vector_index = VectorStoreIndex(
+        nodes=nodes, 
+        embed_model=OpenAIEmbedding(model_name="text-embedding-ada-002")
+    )
+    vector_retriever = vector_index.as_retriever(similarity_top_k=top_k)
+    bm25_retriever = BM25Retriever.from_defaults(nodes=nodes, similarity_top_k=top_k)
+
+    retriever = QueryFusionRetriever(
+        [vector_retriever, bm25_retriever], 
+        similarity_top_k=top_k, 
+        num_queries=3, 
+        mode="dist_based_score",
+        llm=OpenAI(model="gpt-3.5-turbo", temperature=0.1)
+    )
+
+    def _retriever(query: str) -> List[NodeWithScore]:
+        return retriever.retrieve(query)
+
+    return _retriever
+
+
+
+class LlamaRetriever(Retriever):
+    def __init__(
+        self, 
+        name: str, 
+        create_retriever_fn: Callable[[List[Document], int], Callable[[str], List[NodeWithScore]]], 
+        search_top_k: int = 1000,
+    ) -> None:
+        self.name = name
+        self.search_top_k
+        self.create_retriever_fn = create_retriever_fn
+
+    def __str__(self):
+        return self.name
+
+    def __call__(
+        self,
+        corpus: Dict[str, Dict[str, Any]],
+        queries: Dict[str, str],
+        **kwargs,
+    ) -> Dict[str, Dict[str, float]]:
+        """
+        Retrieve relevant documents for each query
+        """
+        
+        documents = []
+        for doc_id, doc in corpus.items():
+            text = doc.pop("text")
+            assert text is not None
+            documents.append(Document(id_=doc_id, text=text, metadata={**doc}))
+        
+        retriever = self.create_retriever_fn(documents)
+
+        query_ids = list(queries.keys())
+        results = {qid: {} for qid in query_ids}
+        for qid in query_ids:
+            query = queries[qid]
+            if isinstance(query, list):
+                # take from mteb: 
+                # https://github.com/embeddings-benchmark/mteb/blob/main/mteb/evaluation/evaluators/RetrievalEvaluator.py#L403
+                query = "; ".join(query)
+            
+            nodes = retriever(query)
+            for node in nodes:
+                # ref_doc_id should point to corpus document id
+                results[qid][node.node.ref_doc_id] = node.score
+        
+        return results
+        
+
+
+retriever = LlamaRetriever("vector_bm25_fusion", create_retriever_fn)
+
+evaluation = AIRBench(
+    benchmark_version="AIR-Bench_24.04",
+    task_types=["long-doc"],      # remove this line if you want to evaluate on all tasks
+    domains=["arxiv"],            # remove this line if you want to evaluate on all domains
+    languages=["en"],             # remove this line if you want to evaluate on all languages
+    # cache_dir="~/.air_bench/"     # path to the cache directory (**NEED ~52GB FOR FULL BENCHMARK**)
+)
+
+evaluation.run(
+    retriever,
+    output_dir="./llama_results",   # path to the output directory, default is "./search_results"
+    overwrite=True             # set to True if you want to overwrite the existing results
+)