File size: 2,256 Bytes
0703e2b | 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 | """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)
# LangChain-style document
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]
# ------------------------------------------------------------------ text utils
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)
|