Create data.py

rag/data.py

@@ -0,0 +1,60 @@
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import List, Optional
+
+from rag.logging_utils import get_logger
+
+logger = get_logger(__name__)
+
+@dataclass
+class Doc:
+    doc_name: str
+    image_path: str
+    text: str
+    company: str  # "Apple" | "Microsoft" | "Unknown"
+
+def infer_company(doc_name: str) -> str:
+    dn = (doc_name or "").lower()
+    if "microsoft" in dn or "msft" in dn:
+        return "Microsoft"
+    if "apple" in dn or "aapl" in dn:
+        return "Apple"
+    return "Unknown"
+
+def load_dataset(dataset_path: str) -> List[Doc]:
+    path = Path(dataset_path)
+    if not path.exists():
+        logger.warning("Dataset file not found: %s", dataset_path)
+        return []
+
+    with path.open("r", encoding="utf-8") as f:
+        raw = json.load(f)
+
+    docs: List[Doc] = []
+    for item in raw:
+        doc_name = item.get("doc_name", "Unknown Document")
+        image_path = item.get("image_path", "")
+        text = (item.get("text") or "").strip()
+        company = item.get("company") or infer_company(doc_name)
+
+        if not text:
+            continue
+
+        docs.append(
+            Doc(
+                doc_name=doc_name,
+                image_path=image_path,
+                text=text,
+                company=company,
+            )
+        )
+
+    logger.info("Loaded %d docs", len(docs))
+    return docs
+
+def filter_docs_by_companies(docs: List[Doc], companies: Optional[List[str]]) -> List[Doc]:
+    if not companies:
+        return docs
+    allowed = set(companies)
+    return [d for d in docs if d.company in allowed]