init

README.md

@@ -1,7 +1,7 @@
 ---
-title: Scifact Prompting
+title: Retrieval Prompting
 emoji: ⚡
-colorFrom: green
+colorFrom: yellow
 colorTo: indigo
 sdk: gradio
 sdk_version: 4.41.0

app.py

@@ -0,0 +1,183 @@
+import gradio as gr
+import pickle
+import numpy as np
+import glob
+from tqdm import tqdm
+import torch
+import torch.nn.functional as F
+from transformers import AutoTokenizer, AutoModel
+from peft import PeftModel
+from tevatron.retriever.searcher import FaissFlatSearcher
+import logging
+import os
+import json
+import spaces
+import ir_datasets
+import subprocess
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Global variables
+CUR_MODEL = "orionweller/repllama-instruct-hard-positives-v2-joint"
+base_model = "meta-llama/Llama-2-7b-hf"
+tokenizer = None
+model = None
+retriever = None
+corpus_lookup = None
+queries = None
+q_lookup = None
+
+def load_model():
+    global tokenizer, model
+    tokenizer = AutoTokenizer.from_pretrained(base_model)
+    tokenizer.pad_token_id = tokenizer.eos_token_id
+    tokenizer.pad_token = tokenizer.eos_token
+    tokenizer.padding_side = "right"
+    
+    base_model_instance = AutoModel.from_pretrained("meta-llama/Llama-2-7b-hf")
+    model = PeftModel.from_pretrained(base_model_instance, CUR_MODEL)
+    model = model.merge_and_unload()
+    model.eval()
+    model.cuda()
+
+def load_corpus_embeddings(dataset_name):
+    global retriever, corpus_lookup
+    corpus_path = f"{dataset_name}/corpus_emb*"
+    index_files = glob.glob(corpus_path)
+    logger.info(f'Pattern match found {len(index_files)} files; loading them into index.')
+
+    p_reps_0, p_lookup_0 = pickle_load(index_files[0])
+    retriever = FaissFlatSearcher(p_reps_0)
+
+    shards = [(p_reps_0, p_lookup_0)] + [pickle_load(f) for f in index_files[1:]]
+    corpus_lookup = []
+    
+    for p_reps, p_lookup in tqdm(shards, desc='Loading shards into index', total=len(index_files)):
+        retriever.add(p_reps)
+        corpus_lookup += p_lookup
+
+def pickle_load(path):
+    with open(path, 'rb') as f:
+        reps, lookup = pickle.load(f)
+    return np.array(reps), lookup
+
+def load_queries(dataset_name):
+    global queries, q_lookup
+    dataset = ir_datasets.load(f"beir/{dataset_name.lower()}/test")
+    
+    queries = []
+    q_lookup = {}
+    for query in dataset.queries_iter():
+        queries.append(query.text)
+        q_lookup[query.query_id] = query.text
+
+def encode_queries(prefix, postfix):
+    global queries
+    input_texts = [f"{prefix}Query:  {query} {postfix}".strip() for query in queries]
+    
+    encoded_embeds = []
+    batch_size = 32  # Adjust as needed
+    
+    for start_idx in range(0, len(input_texts), batch_size):
+        batch_input_texts = input_texts[start_idx: start_idx + batch_size]
+        
+        inputs = tokenizer(batch_input_texts, padding=True, truncation=True, return_tensors="pt").to(model.device)
+        
+        with torch.no_grad():
+            outputs = model(**inputs)
+            embeds = outputs.last_hidden_state[:, 0, :]  # Use [CLS] token embedding
+            embeds = F.normalize(embeds, p=2, dim=-1)
+            encoded_embeds.append(embeds.cpu().numpy())
+    
+    return np.concatenate(encoded_embeds, axis=0)
+
+def search_queries(q_reps, depth=1000):
+    all_scores, all_indices = retriever.search(q_reps, depth)
+    psg_indices = [[str(corpus_lookup[x]) for x in q_dd] for q_dd in all_indices]
+    return all_scores, np.array(psg_indices)
+
+def write_ranking(corpus_indices, corpus_scores, ranking_save_file):
+    with open(ranking_save_file, 'w') as f:
+        for qid, q_doc_scores, q_doc_indices in zip(q_lookup.keys(), corpus_scores, corpus_indices):
+            score_list = [(s, idx) for s, idx in zip(q_doc_scores, q_doc_indices)]
+            score_list = sorted(score_list, key=lambda x: x[0], reverse=True)
+            for rank, (s, idx) in enumerate(score_list, 1):
+                f.write(f'{qid} Q0 {idx} {rank} {s} pyserini\n')
+
+def evaluate_with_subprocess(dataset, ranking_file):
+    # Convert to TREC format
+    trec_file = f"rank.{dataset}.trec"
+    convert_cmd = [
+        "python", "-m", "tevatron.utils.format.convert_result_to_trec",
+        "--input", ranking_file,
+        "--output", trec_file,
+        "--remove_query"
+    ]
+    subprocess.run(convert_cmd, check=True)
+
+    # Evaluate using trec_eval
+    eval_cmd = [
+        "python", "-m", "pyserini.eval.trec_eval",
+        "-c", "-mrecall.100", "-mndcg_cut.10",
+        f"beir-v1.0.0-{dataset}-test", trec_file
+    ]
+    result = subprocess.run(eval_cmd, capture_output=True, text=True, check=True)
+
+    # Parse the output
+    lines = result.stdout.strip().split('\n')
+    ndcg_10 = float(lines[0].split()[-1])
+    recall_100 = float(lines[1].split()[-1])
+
+    # Clean up temporary files
+    os.remove(ranking_file)
+    os.remove(trec_file)
+
+    return f"nDCG@10: {ndcg_10:.4f}, Recall@100: {recall_100:.4f}"
+
+@spaces.GPU
+def run_evaluation(dataset, prefix, postfix):
+    global queries, q_lookup
+    
+    # Load corpus embeddings and queries if not already loaded
+    if retriever is None or queries is None:
+        load_corpus_embeddings(dataset)
+        load_queries(dataset)
+    
+    # Encode queries
+    q_reps = encode_queries(prefix, postfix)
+    
+    # Search
+    all_scores, psg_indices = search_queries(q_reps)
+    
+    # Write ranking
+    ranking_file = f"temp_ranking_{dataset}.txt"
+    write_ranking(psg_indices, all_scores, ranking_file)
+    
+    # Evaluate
+    results = evaluate_with_subprocess(dataset, ranking_file)
+    
+    return results
+
+def gradio_interface(dataset, prefix, postfix):
+    return run_evaluation(dataset, prefix, postfix)
+
+# Load model
+load_model()
+
+# Create Gradio interface
+iface = gr.Interface(
+    fn=gradio_interface,
+    inputs=[
+        gr.Dropdown(choices=["scifact", "arguana"], label="Dataset"),
+        gr.Textbox(label="Prefix prompt"),
+        gr.Textbox(label="Postfix prompt")
+    ],
+    outputs=gr.Textbox(label="Evaluation Results"),
+    title="Query Evaluation with Custom Prompts",
+    description="Select a dataset and enter prefix and postfix prompts to evaluate queries using Pyserini."
+)
+
+# Launch the interface
+iface.launch()

arguana/corpus_emb.0.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:21742104cd3b5ff805fe1a7432c960b6933159c1092d49f5c4cad74922916a9b
+size 35619068

arguana/corpus_emb.1.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:37b86d66f10459bc2c032741f70e9f1edba68b2b3fbf9c7b1c1bb6f6139fef02
+size 35619074

arguana/corpus_emb.2.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f0ebf48b7b6ad18402a90da0ca92b71128b43d81a85fdb425edcaa3767f213d1
+size 35602692

arguana/corpus_emb.3.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:82bb37e7f6d7e0728da781aaccaad49830708cd44abc455f8c3a81db1e4b4b0f
+size 35602679

packages.txt

@@ -0,0 +1 @@
+default-jre

requirements.txt

@@ -0,0 +1,8 @@
+gradio==4.39.0
+pyserini==0.23.0
+faiss-cpu==1.7.4
+torch==2.1.0
+ir_datasets
+peft==0.12.0
+ir_datasets==0.5.8
+tevatron @ git+https://github.com/texttron/tevatron@7d298b4

scifact/corpus_emb.0.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0bb98e68350983519732b0b39e8f98ec0225abd2c68775e7317da9b17f0db1dd
+size 21247618

scifact/corpus_emb.1.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3dd3501342754aeb2ffb895480868e0976895bded3e5accbd8e5b6fa404e5484
+size 21247619

scifact/corpus_emb.2.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0e1a98c698cbe367bc1abc789da76794a8e79e92743059b26faafbd34808aa15
+size 21247619

scifact/corpus_emb.3.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:911c8d6654bfb14a3d68363c96a70462348cfbbf35a591e020877ed28591339c
+size 21231225