"""Ask tiny-search-reader-60M a question about a list of search results. build_prompt(tokenizer, question, snippets) -> prompt token ids that fit the 512-token context answer_transformers(question, snippets) -> answer string, using the Hugging Face files in this folder answer_gguf(question, snippets) -> answer string, using a GGUF file through llama-cpp-python A snippet is a dict {"title": "...", "text": "..."} (a plain string is read as text with no title). Results go in the order the search engine returned them. The answer is a few words, or "not found" when the results do not answer the question. Needs: `tokenizers` for build_prompt; `torch` + `transformers` for answer_transformers; `tokenizers` + `llama-cpp-python` for answer_gguf. Run `python searchreader.py` for a demo. """ import atexit import os import re HERE = os.path.dirname(os.path.abspath(__file__)) TOKENIZER_JSON = os.path.join(HERE, "tokenizer.json") DEFAULT_GGUF = os.path.join(HERE, "gguf", "tiny-search-reader-60M-Q8_0.gguf") CTX = 512 # the model's context length MAX_NEW = 24 # answer tokens; the exams were graded with 24 NOT_FOUND = "not found" QUESTION, SNIPPET, ANSWER, EOT, PAD = "<|question|>", "<|snippet|>", "<|answer|>", "<|endoftext|>", "<|pad|>" EOT_ID, PAD_ID = 0, 4 # ---- prompt building (a port of searchfmt.clean / render_one / Prompt.build as used at test time) ---- def clean(t): return re.sub(r"\s+", " ", (t or "").replace("\u00a0", " ")).strip() def render_one(i, s): if isinstance(s, str): s = {"text": s} title, text = clean(s.get("title")), clean(s.get("text")) return f"[{i}] {title}: {text}" if title else f"[{i}] {text}" class _Tok: """The same three calls on a `tokenizers.Tokenizer` or on a transformers tokenizer (AutoTokenizer).""" def __init__(self, tok): self.tok = tok self.raw = hasattr(tok, "token_to_id") # tokenizers.Tokenizer def encode(self, text): if self.raw: return self.tok.encode(text).ids return self.tok.encode(text, add_special_tokens=False, verbose=False) # no "longer than 512" warning def token_to_id(self, t): return self.tok.token_to_id(t) if self.raw else self.tok.convert_tokens_to_ids(t) def decode(self, ids): if self.raw: return self.tok.decode(ids) return self.tok.decode(ids, skip_special_tokens=True) def _build(tok, question, snippets, ctx, reserve): q, s, a = (tok.token_to_id(t) for t in (QUESTION, SNIPPET, ANSWER)) head = [q] + tok.encode(" " + clean(question)) + [s] room = ctx - len(head) - 1 - reserve parts = [tok.encode(" " + render_one(i + 1, sn)) for i, sn in enumerate(snippets)] keep = list(range(len(parts))) # Too long: drop results from the end, one at a time, but always keep the first. while sum(len(parts[i]) for i in keep) > room and len(keep) > 1: keep.pop() ids = [t for i in keep for t in parts[i]] # Still too long (one huge first result): cut it at the token limit. if len(ids) > room: ids = ids[:room] return head + ids + [a] def build_prompt(tokenizer, question, snippets, ctx=CTX, reserve=MAX_NEW): """Token ids of "<|question|> {question}<|snippet|> [1] {title}: {text} [2] ...<|answer|>". Results keep their order. If they do not fit in ctx - reserve tokens, results are dropped from the end until they do. The first one always stays, and if it alone is still too long it is cut short. A question too long to leave any room is cut to its first ctx // 2 tokens. This is exactly what the model was graded on (finetune_search.Exam.prompt); at inference no gold answer is known, so the training-time rule in searchfmt.Prompt.build that protects snippets holding the answer never applies. tokenizer: a `tokenizers.Tokenizer` (Tokenizer.from_file("tokenizer.json")) or a transformers tokenizer.""" tok = _Tok(tokenizer) ids = _build(tok, question, snippets, ctx, reserve) if len(ids) > ctx - reserve: q = tok.encode(" " + clean(question)) ids = _build(tok, tok.decode(q[:ctx // 2]), snippets, ctx, reserve) return ids # ---- answering ---- _cache = {} @atexit.register def close(): """Free cached llama.cpp models. Runs at exit too: letting Python tear them down late can abort on macOS.""" for key in [k for k in _cache if k[0] == "gguf"]: _cache.pop(key)[0].close() def _hf(model_dir): if ("hf", model_dir) not in _cache: from transformers import AutoModelForCausalLM, AutoTokenizer tok = AutoTokenizer.from_pretrained(model_dir) model = AutoModelForCausalLM.from_pretrained(model_dir).eval() # fp32, as stored _cache["hf", model_dir] = model, tok return _cache["hf", model_dir] def answer_transformers(question, snippets, model_dir=HERE, max_new_tokens=MAX_NEW): """Greedy answer with transformers (fp32 on CPU unless you move the model).""" import torch model, tok = _hf(model_dir) ids = build_prompt(tok, question, snippets, reserve=max_new_tokens) x = torch.tensor([ids], device=model.device) with torch.no_grad(): out = model.generate(x, attention_mask=torch.ones_like(x), max_new_tokens=max_new_tokens, do_sample=False, eos_token_id=EOT_ID, pad_token_id=PAD_ID) new = out[0, len(ids):].tolist() if EOT_ID in new: new = new[:new.index(EOT_ID)] return tok.decode(new, skip_special_tokens=True).strip() def _gguf(gguf_path, tokenizer_path, n_threads): key = ("gguf", gguf_path, tokenizer_path, n_threads) if key not in _cache: from llama_cpp import Llama from tokenizers import Tokenizer llm = Llama(model_path=gguf_path, n_ctx=CTX, n_batch=CTX, n_threads=n_threads, n_gpu_layers=0, verbose=False) _cache[key] = llm, Tokenizer.from_file(tokenizer_path) return _cache[key] def answer_gguf(question, snippets, gguf_path=DEFAULT_GGUF, tokenizer_path=TOKENIZER_JSON, n_threads=4, max_new_tokens=MAX_NEW): """Greedy answer with llama-cpp-python on CPU. The prompt goes in as token ids from our own tokenizer, so the special tokens are single ids and no BOS is added. Decoding stops at <|endoftext|> (id 0) or after max_new_tokens.""" llm, tok = _gguf(gguf_path, tokenizer_path, n_threads) ids = build_prompt(tok, question, snippets, reserve=max_new_tokens) llm.reset() # evaluate every prompt from scratch out = [] for t in llm.generate(ids, temp=0.0, top_k=1, top_p=1.0, min_p=0.0, repeat_penalty=1.0, reset=True): if t == EOT_ID: break out.append(t) if len(out) == max_new_tokens: break return tok.decode(out).strip() # ---- demo (made-up results) ---- DEMO = [ ("When does the Harbor Street library open on Saturday?", [ {"title": "Harbor Street Library - Hours and location", "text": "The Harbor Street Library is open Monday to Friday from 8 am to 8 pm. On Saturday the library " "opens at 10 am and closes at 4 pm. It is closed on Sunday."}, {"title": "Harbor Street Library events", "text": "Story time for children runs every Wednesday morning in the reading room."}, {"title": "City libraries | Parking", "text": "Free parking is available behind the Harbor Street Library for up to two hours."}, ]), ("Who won the 2031 Harbor Street Chess Open?", [ {"title": "Harbor Street Chess Open 2025 results", "text": "Mara Lindqvist won the 2025 Harbor Street Chess Open with 7 points from 9 games."}, {"title": "Harbor Street Chess Open 2026", "text": "The 2026 Harbor Street Chess Open was won by Tomas Varga after a tie-break against Mara Lindqvist."}, {"title": "About the Harbor Street Chess Open", "text": "The open has been played every June at the Harbor Street Library since 2019."}, ]), ] if __name__ == "__main__": import argparse ap = argparse.ArgumentParser(description="Demo: answer two made-up questions from made-up search results.") ap.add_argument("--backend", choices=["transformers", "gguf", "both"], default="both") ap.add_argument("--gguf", default=DEFAULT_GGUF, help="GGUF file for the llama-cpp-python backend") a = ap.parse_args() for question, snippets in DEMO: print("Q:", question) if a.backend in ("transformers", "both"): print(" transformers:", answer_transformers(question, snippets)) if a.backend in ("gguf", "both"): print(f" gguf ({os.path.basename(a.gguf)}):", answer_gguf(question, snippets, gguf_path=a.gguf))