working

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 import pickle
 import numpy as np
 import glob
-from tqdm import tqdm
+import tqdm
 import torch
 import torch.nn.functional as F
 from transformers import AutoTokenizer, AutoModel
@@ -13,50 +13,81 @@ import os
 import json
 import spaces
 import ir_datasets
-import subprocess
+import pytrec_eval
+from huggingface_hub import login
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+# Authenticate with HF_TOKEN
+login(token=os.environ['HF_TOKEN'])
+
 # Global variables
 CUR_MODEL = "orionweller/repllama-instruct-hard-positives-v2-joint"
-base_model = "meta-llama/Llama-2-7b-hf"
+BASE_MODEL = "meta-llama/Llama-2-7b-hf"
 tokenizer = None
 model = None
-retriever = None
-corpus_lookup = None
-queries = None
-q_lookup = None
+retrievers = {}
+corpus_lookups = {}
+queries = {}
+q_lookups = {}
+qrels = {}
+datasets = ["scifact", "arguana"]
+current_dataset = "scifact"
+
+def pool(last_hidden_states, attention_mask):
+    last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
+    sequence_lengths = attention_mask.sum(dim=1) - 1
+    batch_size = last_hidden.shape[0]
+    return last_hidden[torch.arange(batch_size, device=last_hidden.device), sequence_lengths]
+
+def create_batch_dict(tokenizer, input_texts, max_length=512):
+    batch_dict = tokenizer(
+        input_texts,
+        max_length=max_length - 1,
+        return_token_type_ids=False,
+        return_attention_mask=False,
+        padding=False,
+        truncation=True
+    )
+    batch_dict['input_ids'] = [input_ids + [tokenizer.eos_token_id] for input_ids in batch_dict['input_ids']]
+    return tokenizer.pad(
+        batch_dict,
+        padding=True,
+        pad_to_multiple_of=8,
+        return_attention_mask=True,
+        return_tensors="pt",
+    )
 
 def load_model():
     global tokenizer, model
-    tokenizer = AutoTokenizer.from_pretrained(base_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")
+    base_model_instance = AutoModel.from_pretrained(BASE_MODEL)
     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*"
+    global retrievers, corpus_lookups
+    corpus_path = f"{dataset_name}/corpus_emb.*.pkl"
     index_files = glob.glob(corpus_path)
-    logger.info(f'Pattern match found {len(index_files)} files; loading them into index.')
+    logger.info(f'Loading {len(index_files)} files into index for {dataset_name}.')
 
     p_reps_0, p_lookup_0 = pickle_load(index_files[0])
-    retriever = FaissFlatSearcher(p_reps_0)
+    retrievers[dataset_name] = FaissFlatSearcher(p_reps_0)
 
     shards = [(p_reps_0, p_lookup_0)] + [pickle_load(f) for f in index_files[1:]]
-    corpus_lookup = []
+    corpus_lookups[dataset_name] = []
     
-    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
+    for p_reps, p_lookup in tqdm.tqdm(shards, desc=f'Loading shards into index for {dataset_name}', total=len(index_files)):
+        retrievers[dataset_name].add(p_reps)
+        corpus_lookups[dataset_name] += p_lookup
 
 def pickle_load(path):
     with open(path, 'rb') as f:
@@ -64,119 +95,104 @@ def pickle_load(path):
     return np.array(reps), lookup
 
 def load_queries(dataset_name):
-    global queries, q_lookup
-    dataset = ir_datasets.load(f"beir/{dataset_name.lower()}/test")
+    global queries, q_lookups, qrels
+    dataset = ir_datasets.load(f"beir/{dataset_name.lower()}" + ("/test" if dataset_name == "scifact" else ""))
     
-    queries = []
-    q_lookup = {}
+    queries[dataset_name] = []
+    q_lookups[dataset_name] = {}
+    qrels[dataset_name] = {}
     for query in dataset.queries_iter():
-        queries.append(query.text)
-        q_lookup[query.query_id] = query.text
+        queries[dataset_name].append(query.text)
+        q_lookups[dataset_name][query.query_id] = query.text
+    
+    for qrel in dataset.qrels_iter():
+        if qrel.query_id not in qrels[dataset_name]:
+            qrels[dataset_name][qrel.query_id] = {}
+        qrels[dataset_name][qrel.query_id][qrel.doc_id] = qrel.relevance
 
-def encode_queries(prefix, postfix):
-    global queries
-    input_texts = [f"{prefix}Query:  {query} {postfix}".strip() for query in queries]
+@spaces.GPU
+def encode_queries(dataset_name, postfix):
+    global queries, tokenizer, model
+    model = model.cuda()
+    input_texts = [f"query: {query.strip()} {postfix}".strip() for query in queries[dataset_name]]
     
     encoded_embeds = []
-    batch_size = 32  # Adjust as needed
+    batch_size = 32
     
-    for start_idx in range(0, len(input_texts), batch_size):
+    for start_idx in tqdm.tqdm(range(0, len(input_texts), batch_size), desc="Encoding queries"):
         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
+        batch_dict = create_batch_dict(tokenizer, batch_input_texts)
+        batch_dict = {k: v.to(model.device) for k, v in batch_dict.items()}
+
+        with torch.cuda.amp.autocast():
+            outputs = model(**batch_dict)
+            embeds = pool(outputs.last_hidden_state, batch_dict['attention_mask'])
             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]
+def search_queries(dataset_name, q_reps, depth=1000):
+    all_scores, all_indices = retrievers[dataset_name].search(q_reps, depth)
+    psg_indices = [[str(corpus_lookups[dataset_name][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}"
+def evaluate(qrels, results, k_values):
+    evaluator = pytrec_eval.RelevanceEvaluator(
+        qrels, {f"ndcg_cut.{k}" for k in k_values} | {f"recall.{k}" for k in k_values}
+    )
+    scores = evaluator.evaluate(results)
 
-@spaces.GPU
-def run_evaluation(dataset, prefix, postfix):
-    global queries, q_lookup
+    metrics = {}
+    for k in k_values:
+        metrics[f"NDCG@{k}"] = round(np.mean([query_scores[f"ndcg_cut_{k}"] for query_scores in scores.values()]), 3)
+        metrics[f"Recall@{k}"] = round(np.mean([query_scores[f"recall_{k}"] for query_scores in scores.values()]), 3)
+
+    return metrics
+
+def run_evaluation(dataset, postfix):
+    global current_dataset
     
-    # Load corpus embeddings and queries if not already loaded
-    if retriever is None or queries is None:
+    if dataset not in retrievers or dataset not in queries:
         load_corpus_embeddings(dataset)
         load_queries(dataset)
     
-    # Encode queries
-    q_reps = encode_queries(prefix, postfix)
+    current_dataset = dataset
     
-    # Search
-    all_scores, psg_indices = search_queries(q_reps)
+    q_reps = encode_queries(dataset, postfix)
+    all_scores, psg_indices = search_queries(dataset, q_reps)
     
-    # Write ranking
-    ranking_file = f"temp_ranking_{dataset}.txt"
-    write_ranking(psg_indices, all_scores, ranking_file)
+    results = {qid: dict(zip(doc_ids, map(float, scores))) 
+               for qid, scores, doc_ids in zip(q_lookups[dataset].keys(), all_scores, psg_indices)}
     
-    # Evaluate
-    results = evaluate_with_subprocess(dataset, ranking_file)
+    metrics = evaluate(qrels[dataset], results, k_values=[10, 100])
     
-    return results
+    return {
+        "NDCG@10": metrics["NDCG@10"],
+        "Recall@100": metrics["Recall@100"]
+    }
 
-def gradio_interface(dataset, prefix, postfix):
-    return run_evaluation(dataset, prefix, postfix)
+def gradio_interface(dataset, postfix):
+    return run_evaluation(dataset, postfix)
 
-# Load model
+# Load model and initial datasets
 load_model()
+for dataset in datasets:
+    print(f"Loading dataset: {dataset}")
+    load_corpus_embeddings(dataset)
+    load_queries(dataset)
 
 # 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")
+        gr.Dropdown(choices=datasets, label="Dataset", value="scifact"),
+        gr.Textbox(label="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."
+    outputs=gr.JSON(label="Evaluation Results"),
+    title="Promptriever Demo",
+    description="Select a dataset and enter a postfix prompt to evaluate the model's performance. Note: it takes about **ten seconds** for each dataset."
 )
 
 # Launch the interface

requirements.txt

@@ -1,8 +1,9 @@
-gradio==4.39.0
+gradio==4.43.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
+pytrec_eval==0.5
 tevatron @ git+https://github.com/texttron/tevatron@7d298b4