| """Documents and text utilities used by the RAG pipeline.""" |
| from __future__ import annotations |
|
|
| import hashlib |
| import re |
| import string |
| from dataclasses import dataclass, field |
| from typing import Any, Dict, List, Union |
|
|
|
|
| @dataclass |
| class Document: |
| """A retrieved passage.""" |
|
|
| text: str |
| title: str = "" |
| id: str = "" |
| metadata: Dict[str, Any] = field(default_factory=dict) |
|
|
| def __post_init__(self) -> None: |
| if not self.id: |
| self.id = hashlib.sha1((self.title + "\x00" + self.text).encode("utf-8")).hexdigest()[:12] |
|
|
| def to_dict(self) -> Dict[str, Any]: |
| return {"id": self.id, "title": self.title, "text": self.text, "metadata": dict(self.metadata)} |
|
|
|
|
| DocLike = Union[Document, str, Dict[str, Any]] |
|
|
|
|
| def as_document(x: DocLike) -> Document: |
| if isinstance(x, Document): |
| return x |
| if isinstance(x, str): |
| return Document(text=x) |
| if isinstance(x, dict): |
| meta = {k: v for k, v in x.items() if k not in ("text", "title", "id", "page_content")} |
| return Document(text=x.get("text", x.get("page_content", "")), title=x.get("title", ""), |
| id=str(x.get("id", "")), metadata=meta) |
| |
| text = getattr(x, "page_content", None) |
| if text is not None: |
| meta = dict(getattr(x, "metadata", {}) or {}) |
| return Document(text=text, title=str(meta.get("title", "")), id=str(meta.get("id", "") or getattr(x, "id", "") or ""), |
| metadata=meta) |
| raise TypeError(f"cannot convert {type(x).__name__} to Document") |
|
|
|
|
| def as_documents(xs) -> List[Document]: |
| return [as_document(x) for x in xs] |
|
|
|
|
| |
| def normalize_answer(s: str) -> str: |
| """SQuAD normalisation: lower case, strip punctuation and articles, squeeze spaces.""" |
| s = s.lower() |
| s = "".join(ch for ch in s if ch not in set(string.punctuation)) |
| s = re.sub(r"\b(a|an|the)\b", " ", s) |
| return " ".join(s.split()) |
|
|
|
|
| def contains_answer(passage: str, answer: str) -> bool: |
| """Grounding check: the normalised answer occurs in the normalised passage.""" |
| a = normalize_answer(answer) |
| return bool(a) and a in normalize_answer(passage) |
|
|
