Spaces:
Sleeping
Sleeping
File size: 1,546 Bytes
c56ab56 |
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 |
from dataclasses import asdict
import json
from typing import Iterator
from dateutil.parser import parse
from langchain.docstore.document import Document
from langchain.document_loaders.base import BaseLoader
from model import Doc
def date_to_int(dt_str: str) -> int:
dt = parse(dt_str)
return int(dt.timestamp())
def get_contents(project_name: str, filename: str) -> Iterator[tuple[Doc, str]]:
"""filename for file with ndjson
{"title": <page title>, "body": <page body>, "id": <page_id>, "ctime": ..., "user": <name>, "url": "https:..."}
{"title": ...}
"""
with open(filename, "r") as f:
obj = [json.loads(line) for line in f]
for data in obj:
title = data["title"]
body = data["body"]
doc = Doc(
project_name=project_name,
id=data["id"],
title=title,
created_at=date_to_int(data["ctime"]),
user=data["user"],
url=data["url"],
)
text = title
if body:
text += "\n\n" + body
yield doc, text
class DocLoader(BaseLoader):
def __init__(self, project_name: str, filename: str):
self.project_name = project_name
self.filename = filename
def lazy_load(self) -> Iterator[Document]:
for doc, text in get_contents(self.project_name, self.filename):
metadata = asdict(doc)
yield Document(page_content=text, metadata=metadata)
def load(self) -> list[Document]:
return list(self.lazy_load())
|