|
|
|
|
|
"""Utilities for loading the ZeroSearch simulation model and performing simulated searches.""" |
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline |
|
|
import functools |
|
|
|
|
|
MODEL_NAME = "sunhaonlp/SearchSimulation_14B" |
|
|
|
|
|
@functools.lru_cache(maxsize=1) |
|
|
def _load_search_pipe(): |
|
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) |
|
|
model = AutoModelForCausalLM.from_pretrained( |
|
|
MODEL_NAME, |
|
|
trust_remote_code=True, |
|
|
device_map="auto" |
|
|
) |
|
|
return pipeline( |
|
|
"text-generation", |
|
|
model=model, |
|
|
tokenizer=tokenizer, |
|
|
max_new_tokens=512, |
|
|
do_sample=False, |
|
|
temperature=0.0, |
|
|
) |
|
|
|
|
|
def simulate_search(query: str, k: int = 5): |
|
|
"""Generate *k* synthetic documents for *query*.""" |
|
|
pipe = _load_search_pipe() |
|
|
prompt = f"SearchSimulation:\nQuery: {query}\nDocuments:" |
|
|
outputs = pipe(prompt, num_return_sequences=k) |
|
|
docs = [] |
|
|
for o in outputs: |
|
|
text = o["generated_text"] |
|
|
docs.append(text.split("Documents:")[-1].strip()) |
|
|
return docs |
|
|
|
