implementing the UNIMPLEMENTED_PARSERS

.gitignore

@@ -123,4 +123,6 @@ carado.moe/
 *.epub
 
 credentials.json
-data/raw/
+data/raw/
+
+*.log

align_data/analysis/analyse_jsonl_data.py

@@ -1,8 +1,9 @@
 from datetime import datetime
 from pathlib import Path
+from collections import defaultdict
+
 import jsonlines
 
-from collections import defaultdict
 
 
 def is_valid_date_format(data_dict, format="%Y-%m-%dT%H:%M:%SZ"):

align_data/analysis/count_tokens.py

@@ -1,3 +1,5 @@
+from typing import Tuple
+
 from transformers import AutoTokenizer
 import jsonlines
 import logging

align_data/common/alignment_dataset.py

@@ -4,30 +4,21 @@
 import time
 from dataclasses import dataclass, field, KW_ONLY
 from pathlib import Path
-from typing import Iterable, List, Optional, Set
+from typing import List, Optional, Dict, Any, Set, Generator, Iterable
+import pytz
+from datetime import datetime
+
 from sqlalchemy import select
 from sqlalchemy.exc import IntegrityError
-from sqlalchemy.orm import joinedload
-
+from sqlalchemy.orm import Query, joinedload
+from sqlalchemy.engine import Result
 import jsonlines
-import pytz
 from dateutil.parser import parse, ParserError
 from tqdm import tqdm
+
 from align_data.db.models import Article, Summary
 from align_data.db.session import make_session
-
-INIT_DICT = {
-    "source": None,
-    "id": None,
-    "text": None,
-    "date_published": None,
-    "title": None,
-    "url": None,
-    "authors": lambda: [],
-    "source_type": None,
-    "status": None,
-    "comments": None,
-}
+from align_data.settings import ARTICLE_MAIN_KEYS
 
 logger = logging.getLogger(__name__)
 
@@ -41,7 +32,16 @@ class AlignmentDataset:
 
     _: KW_ONLY
 
-    files_path = Path("")
+    id_fields: List[str] = field(default_factory=lambda: ["url", "title"])
+    """A list of fields to use as the id of the entry. If not set, will use ['url', 'title']"""
+
+    # Internal housekeeping variables
+    _outputted_items: Set[str] = field(default_factory=set)
+    """A set of the ids of all previously processed items"""
+
+    data_path: Path = field(init=False)
+    raw_data_path: Path = field(init=False)
+    files_path: Path = field(init=False)
     """The path where data can be found. Usually a folder"""
 
     done_key = "id"
@@ -55,23 +55,12 @@ class AlignmentDataset:
     batch_size = 20
     """The number of items to collect before flushing to the database."""
 
-    # Internal housekeeping variables
-    _entry_idx = 0
-    """Used internally for writing debugging info - each file write will increment it"""
-    _outputted_items = set()
-    """A set of the ids of all previously processed items"""
-    _: KW_ONLY
-    id_fields: List[str] = field(default_factory=lambda: ["url", "title"])
-    """A list of fields to use as the id of the entry. If not set, will use ['url', 'title']"""
-
     def __str__(self) -> str:
         return self.name
 
-    def __post_init__(self, data_path=Path(__file__).parent / "../../data/"):
-        self.data_path = data_path
+    def __post_init__(self, data_path: Optional[Path] = None):
+        self.data_path = data_path or (Path(__file__).parent / "../../data/").resolve()
         self.raw_data_path = self.data_path / "raw"
-
-        # set the default place to look for data
         self.files_path = self.raw_data_path / self.name
 
     def _add_authors(self, article: Article, authors: List[str]) -> Article:
@@ -81,33 +70,36 @@ def _add_authors(self, article: Article, authors: List[str]) -> Article:
             article.authors = ",".join(article.authors[:1024].split(",")[:-1])
         return article
 
-    def make_data_entry(self, data, **kwargs) -> Article:
+    def make_data_entry(self, data: Dict[str, Any], **kwargs) -> Article:
         data = dict(data, **kwargs)
         summary = data.pop("summary", None)
         authors = data.pop("authors", [])
 
         article = Article(
             id_fields=self.id_fields,
-            meta={k: v for k, v in data.items() if k not in INIT_DICT and v is not None},
-            **{k: v for k, v in data.items() if k in INIT_DICT},
+            meta={k: v for k, v in data.items() if k not in ARTICLE_MAIN_KEYS and v is not None},
+            **{k: v for k, v in data.items() if k in ARTICLE_MAIN_KEYS},
         )
         self._add_authors(article, authors)
         if summary:
             article.summaries.append(Summary(text=summary, source=self.name))
         return article
 
-    def to_jsonl(self, out_path=None, filename=None) -> Path:
+    def to_jsonl(self, out_path: Path = None, filename: str = None):
         if not out_path:
-            out_path = Path(__file__).parent / "../../data/"
+            out_path=self.data_path
 
-        if not filename:
-            filename = f"{self.name}.jsonl"
-        filename = Path(out_path) / filename
+        filename = filename or f"{self.name}.jsonl"
+        filepath = out_path / filename
 
-        with jsonlines.open(filename, "w") as jsonl_writer:
+        with jsonlines.open(filepath, "w") as jsonl_writer:
             for article in self.read_entries():
                 jsonl_writer.write(article.to_dict())
-        return filename.resolve()
+        return filepath.resolve()
+
+    @property
+    def _query_items(self):
+        return select(Article).where(Article.source == self.name)
 
     @property
     def _query_items(self):
@@ -119,19 +111,22 @@ def read_entries(self, sort_by=None):
             query = self._query_items.options(joinedload(Article.summaries))
             if sort_by is not None:
                 query = query.order_by(sort_by)
-            for item in session.scalars(query).unique():
+
+            result: Result = session.scalars(query)
+            for item in result.unique():
                 yield item
 
     def _add_batch(self, session, batch):
         session.add_all(batch)
 
     def add_entries(self, entries):
-        def commit():
+        def commit() -> bool:
             try:
                 session.commit()
                 return True
             except IntegrityError:
                 session.rollback()
+                return False
 
         with make_session() as session:
             items = iter(entries)
@@ -183,23 +178,24 @@ def unprocessed_items(self, items=None) -> Iterable:
         based on the contents of the output file.
         """
         self.setup()
+        items = items or self.items_list
 
         def not_processed(item):
             # NOTE: `self._outputted_items` reads in all items. Which could potentially be a lot. If this starts to
             # cause problems (e.g. massive RAM usage, big slow downs) then it will have to be switched around, so that
             # this function runs a query to check if the item is in the database rather than first getting all done_keys.
             # If it get's to that level, consider batching it somehow
             return self.get_item_key(item) not in self._outputted_items
-
-        filtered = filter(not_processed, items or self.items_list)
+        
+        items_to_process = filter(not_processed, items)
 
         # greedily fetch all items if not lazy eval. This makes the progress bar look nice
         if not self.lazy_eval:
-            filtered = list(filtered)
+            items_to_process = list(items_to_process)
 
-        return filtered
+        return items_to_process
 
-    def fetch_entries(self):
+    def fetch_entries(self) -> Article:
         """Get all entries to be written to the file."""
         for item in tqdm(self.unprocessed_items(), desc=f"Processing {self.name}"):
             entry = self.process_entry(item)
@@ -216,11 +212,11 @@ def process_entry(self, entry) -> Optional[Article]:
         raise NotImplementedError
 
     @staticmethod
-    def _format_datetime(date) -> str:
+    def _format_datetime(date: datetime) -> str:
         return date.strftime("%Y-%m-%dT%H:%M:%SZ")
 
     @staticmethod
-    def _get_published_date(date) -> Optional[datetime]:
+    def _get_published_date(date: str) -> Optional[datetime]:
         try:
             # Totally ignore any timezone info, forcing everything to UTC
             return parse(str(date)).replace(tzinfo=pytz.UTC)

align_data/common/html_dataset.py

@@ -1,23 +1,22 @@
 import pytz
-import regex as re
 import logging
 from datetime import datetime
-from dateutil.parser import parse
 from dataclasses import dataclass, field, KW_ONLY
 from urllib.parse import urljoin
 from typing import List
+import re
 
 import requests
 import feedparser
 from bs4 import BeautifulSoup
+from bs4.element import ResultSet, Tag
 from markdownify import markdownify
 
 from align_data.common.alignment_dataset import AlignmentDataset
-
 logger = logging.getLogger(__name__)
 
 
-@dataclass
+@dataclass(kw_only=KW_ONLY)
 class HTMLDataset(AlignmentDataset):
     """
     Fetches articles from a different blog by collecting links to articles from an index page.
@@ -27,36 +26,35 @@ class HTMLDataset(AlignmentDataset):
     done_key = "url"
 
     authors: List[str] = field(default_factory=list)
-    _: KW_ONLY
-    source_key: str = None
-    summary_key: str = None
 
     item_selector = "article"
     title_selector = "article h1"
     text_selector = "article"
     source_type = "blog"
     ignored_selectors = []
 
-    def extract_authors(self, article):
+    def extract_authors(self, article): #TODO: make this work
         return self.authors
 
-    def get_item_key(self, item):
-        article_url = item.find_all("a")[0]["href"].split("?")[0]
-        return urljoin(self.url, article_url)
+
+    def get_item_key(self, item: Tag) -> str:
+        first_href = item.find("a")["href"]
+        href_base, *_ = first_href.split("?")
+        return urljoin(self.url, href_base)
 
     @property
-    def items_list(self):
+    def items_list(self) -> ResultSet[Tag]:
         logger.info(f"Fetching entries from {self.url}")
         response = requests.get(self.url, allow_redirects=True)
         soup = BeautifulSoup(response.content, "html.parser")
         articles = soup.select(self.item_selector)
         logger.info(f"Found {len(articles)} articles")
         return articles
 
-    def _extra_values(self, contents):
+    def _extra_values(self, contents: BeautifulSoup):
         return {}
 
-    def process_entry(self, article):
+    def process_entry(self, article: Tag):
         article_url = self.get_item_key(article)
         contents = self._get_contents(article_url)
 
@@ -79,7 +77,7 @@ def process_entry(self, article):
             }
         )
 
-    def _get_contents(self, url):
+    def _get_contents(self, url: str):
         logger.info("Fetching {}".format(url))
         resp = requests.get(url, allow_redirects=True)
         return BeautifulSoup(resp.content, "html.parser")
@@ -97,7 +95,7 @@ def _get_text(self, contents):
 
     def _find_date(self, items):
         for i in items:
-            if re.match("\w+ \d{1,2}, \d{4}", i.text):
+            if re.match(r"\w+ \d{1,2}, \d{4}", i.text):
                 return datetime.strptime(i.text, "%b %d, %Y").replace(tzinfo=pytz.UTC)
 
     def _extract_markdown(self, element):
@@ -132,7 +130,7 @@ def _get_text(self, item):
         text = item.get("content") and item["content"][0].get("value")
         return self._extract_markdown(text)
 
-    def _get_contents(self, url):
+    def _get_contents(self, url: str):
         item = self.items[url]
         if "content" in item:
             return item

align_data/db/session.py

@@ -1,18 +1,18 @@
 from typing import List
 import logging
 
-from contextlib import contextmanager
-from sqlalchemy import create_engine
+from contextlib import contextmanager, AbstractContextManager
+from sqlalchemy import create_engine, Engine
 from sqlalchemy.orm import Session
+
 from align_data.settings import DB_CONNECTION_URI
 from align_data.db.models import Article
 
-
 logger = logging.getLogger(__name__)
 
 
 @contextmanager
-def make_session(auto_commit=False):
+def make_session(auto_commit: bool = False) -> AbstractContextManager[Session]:
     engine = create_engine(DB_CONNECTION_URI, echo=False)
     with Session(engine).no_autoflush as session:
         yield session
@@ -22,6 +22,10 @@ def make_session(auto_commit=False):
 
 def stream_pinecone_updates(session, custom_sources: List[str]):
     """Yield Pinecone entries that require an update."""
-    yield from session.query(Article).filter(
-        Article.pinecone_update_required.is_(True)
-    ).filter(Article.source.in_(custom_sources)).yield_per(1000)
+    yield from (
+        session
+        .query(Article)
+        .filter(Article.pinecone_update_required.is_(True))
+        .filter(Article.source.in_(custom_sources))
+        .yield_per(1000)
+        )

align_data/pinecone/update_pinecone.py

@@ -23,7 +23,6 @@
     EMBEDDING_LENGTH_BIAS,
 )
 
-
 logger = logging.getLogger(__name__)