Spaces:
Running
Running
File size: 3,424 Bytes
8405b55 b96b498 8405b55 b96b498 8405b55 b96b498 8405b55 b96b498 8405b55 b96b498 8405b55 b96b498 8405b55 b96b498 8405b55 b96b498 8405b55 b96b498 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | """Query transformation strategies via LiteLLM, as used by Notebook 6.
`QueryTransformCache` wraps `rewrite_query` / `expand_queries` / `hyde_generate`
with a disk-backed cache (keyed by `(transform_type, question)`), so repeated
demo runs for the same question + strategy don't re-call the LLM.
Sync methods are used by the single-query path (called from a thread via
`asyncio.to_thread`). A `threading.Lock` guards the cache file writes so
concurrent benchmark threads do not corrupt the JSON on disk.
"""
import hashlib
import json
import re
import threading
from pathlib import Path
from src import settings
from src.generation.llm_client import DEFAULT_MODEL, generate
REWRITE_PROMPT = """Rewrite the following question to make it more effective for retrieving \
relevant passages from a document corpus. Preserve the original meaning. Output ONLY \
the rewritten question, nothing else.
Question: {question}
Rewritten question:"""
MULTIQUERY_PROMPT = """Generate {n} alternative phrasings of the following question. Each \
alternative should preserve the original meaning but use different wording or \
emphasis, to help retrieve relevant passages from a document corpus.
Question: {question}
Respond with ONLY a JSON array of {n} strings, e.g. ["...", "..."]"""
HYDE_PROMPT = """Write a short hypothetical passage (2-4 sentences) that would directly \
answer the following question. Write it as if it were an excerpt from a reference \
document, not as a conversational response. Output ONLY the passage, nothing else.
Question: {question}
Hypothetical passage:"""
_write_lock = threading.Lock()
class QueryTransformCache:
def __init__(self, path: Path = None):
self.path = path or (settings.APP_CACHE_DIR / "query_transforms.json")
if self.path.exists():
self._cache = json.loads(self.path.read_text())
else:
self._cache = {}
def _get_or_compute(self, question: str, kind: str, compute_fn):
key = hashlib.md5(f"{kind}|{question}".encode("utf-8")).hexdigest()
if key in self._cache:
return self._cache[key]
result = compute_fn(question)
with _write_lock:
self._cache[key] = result
self.path.write_text(json.dumps(self._cache))
return result
def rewrite_query(self, question: str, model: str = DEFAULT_MODEL) -> str:
def _compute(q):
return generate(REWRITE_PROMPT.format(question=q), model=model).strip() or q
return self._get_or_compute(question, "rewrite", _compute)
def expand_queries(self, question: str, n: int = 2, model: str = DEFAULT_MODEL) -> list:
def _compute(q):
raw = generate(MULTIQUERY_PROMPT.format(question=q, n=n), model=model)
match = re.search(r"\[.*\]", raw, re.DOTALL)
try:
variants = json.loads(match.group(0)) if match else []
except Exception:
variants = []
variants = [v.strip() for v in variants if isinstance(v, str) and v.strip()]
return variants[:n] or [q]
return self._get_or_compute(question, f"multiquery_{n}", _compute)
def hyde_generate(self, question: str, model: str = DEFAULT_MODEL) -> str:
def _compute(q):
return generate(HYDE_PROMPT.format(question=q), model=model).strip() or q
return self._get_or_compute(question, "hyde", _compute)
|