Move preprocessing & registry to separate modules

src/gistillery/ml.py

@@ -1,36 +1,11 @@
 import abc
 from typing import Any
 import logging
-import re
-
-import httpx
-
-from gistillery.base import JobInput
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
 
 
-class Processor(abc.ABC):
-    def get_name(self) -> str:
-        return self.__class__.__name__
-
-    def __call__(self, job: JobInput) -> str:
-        _id = job.id
-        logger.info(f"Processing {input} with {self.__class__.__name__} (id={_id[:8]})")
-        result = self.process(job)
-        logger.info(f"Finished processing input (id={_id[:8]})")
-        return result
-
-    @abc.abstractmethod
-    def process(self, input: JobInput) -> str:
-        raise NotImplementedError
-
-    @abc.abstractmethod
-    def match(self, input: JobInput) -> bool:
-        raise NotImplementedError
-
-
 class Summarizer(abc.ABC):
     def __init__(
         self, model_name: str, model: Any, tokenizer: Any, generation_config: Any
@@ -59,40 +34,6 @@ class Tagger(abc.ABC):
         raise NotImplementedError
 
 
-class MlRegistry:
-    def __init__(self) -> None:
-        self.processors: list[Processor] = []
-        self.summerizer: Summarizer | None = None
-        self.tagger: Tagger | None = None
-        self.model = None
-        self.tokenizer = None
-
-    def register_processor(self, processor: Processor) -> None:
-        self.processors.append(processor)
-
-    def register_summarizer(self, summarizer: Summarizer) -> None:
-        self.summerizer = summarizer
-
-    def register_tagger(self, tagger: Tagger) -> None:
-        self.tagger = tagger
-
-    def get_processor(self, input: JobInput) -> Processor:
-        assert self.processors
-        for processor in self.processors:
-            if processor.match(input):
-                return processor
-
-        return RawTextProcessor()
-
-    def get_summarizer(self) -> Summarizer:
-        assert self.summerizer
-        return self.summerizer
-
-    def get_tagger(self) -> Tagger:
-        assert self.tagger
-        return self.tagger
-
-
 class HfTransformersSummarizer(Summarizer):
     def __init__(
         self, model_name: str, model: Any, tokenizer: Any, generation_config: Any
@@ -152,34 +93,3 @@ class HfTransformersTagger(Tagger):
 
     def get_name(self) -> str:
         return f"{self.__class__.__name__}({self.model_name})"
-
-
-class RawTextProcessor(Processor):
-    def match(self, input: JobInput) -> bool:
-        return True
-
-    def process(self, input: JobInput) -> str:
-        return input.content
-
-
-class DefaultUrlProcessor(Processor):
-    def __init__(self) -> None:
-        self.client = httpx.Client()
-        self.regex = re.compile(r"(https?://[^\s]+)")
-        self.url = None
-        self.template = "{url}\n\n{content}"
-
-    def match(self, input: JobInput) -> bool:
-        urls = list(self.regex.findall(input.content))
-        if len(urls) == 1:
-            self.url = urls[0]
-            return True
-        return False
-
-    def process(self, input: JobInput) -> str:
-        """Get content of website and return it as string"""
-        assert isinstance(self.url, str)
-        text = self.client.get(self.url).text
-        assert isinstance(text, str)
-        text = self.template.format(url=self.url, content=text)
-        return text

src/gistillery/preprocessing.py

@@ -0,0 +1,62 @@
+import abc
+import logging
+import re
+
+import httpx
+
+from gistillery.base import JobInput
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+
+
+
+class Processor(abc.ABC):
+    def get_name(self) -> str:
+        return self.__class__.__name__
+
+    def __call__(self, job: JobInput) -> str:
+        _id = job.id
+        logger.info(f"Processing {input} with {self.__class__.__name__} (id={_id[:8]})")
+        result = self.process(job)
+        logger.info(f"Finished processing input (id={_id[:8]})")
+        return result
+
+    @abc.abstractmethod
+    def process(self, input: JobInput) -> str:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def match(self, input: JobInput) -> bool:
+        raise NotImplementedError
+
+
+class RawTextProcessor(Processor):
+    def match(self, input: JobInput) -> bool:
+        return True
+
+    def process(self, input: JobInput) -> str:
+        return input.content.strip()
+
+
+class DefaultUrlProcessor(Processor):
+    def __init__(self) -> None:
+        self.client = httpx.Client()
+        self.regex = re.compile(r"(https?://[^\s]+)")
+        self.url = None
+        self.template = "{url}\n\n{content}"
+
+    def match(self, input: JobInput) -> bool:
+        urls = list(self.regex.findall(input.content.strip()))
+        if len(urls) == 1:
+            self.url = urls[0]
+            return True
+        return False
+
+    def process(self, input: JobInput) -> str:
+        """Get content of website and return it as string"""
+        assert isinstance(self.url, str)
+        text = self.client.get(self.url).text
+        assert isinstance(text, str)
+        text = self.template.format(url=self.url, content=text)
+        return text

src/gistillery/registry.py

@@ -0,0 +1,38 @@
+from gistillery.ml import Summarizer, Tagger
+from gistillery.preprocessing import Processor, RawTextProcessor
+
+from gistillery.base import JobInput
+
+
+class MlRegistry:
+    def __init__(self) -> None:
+        self.processors: list[Processor] = []
+        self.summerizer: Summarizer | None = None
+        self.tagger: Tagger | None = None
+        self.model = None
+        self.tokenizer = None
+
+    def register_processor(self, processor: Processor) -> None:
+        self.processors.append(processor)
+
+    def register_summarizer(self, summarizer: Summarizer) -> None:
+        self.summerizer = summarizer
+
+    def register_tagger(self, tagger: Tagger) -> None:
+        self.tagger = tagger
+
+    def get_processor(self, input: JobInput) -> Processor:
+        assert self.processors
+        for processor in self.processors:
+            if processor.match(input):
+                return processor
+
+        return RawTextProcessor()
+
+    def get_summarizer(self) -> Summarizer:
+        assert self.summerizer
+        return self.summerizer
+
+    def get_tagger(self) -> Tagger:
+        assert self.tagger
+        return self.tagger

src/gistillery/worker.py

@@ -3,13 +3,9 @@ from dataclasses import dataclass
 
 from gistillery.base import JobInput
 from gistillery.db import get_db_cursor
-from gistillery.ml import (
-    DefaultUrlProcessor,
-    HfTransformersSummarizer,
-    HfTransformersTagger,
-    MlRegistry,
-    RawTextProcessor,
-)
+from gistillery.ml import HfTransformersSummarizer, HfTransformersTagger
+from gistillery.preprocessing import DefaultUrlProcessor, RawTextProcessor
+from gistillery.registry import MlRegistry
 
 SLEEP_INTERVAL = 5
 

tests/test_app.py

@@ -37,7 +37,9 @@ class TestWebservice:
     @pytest.fixture
     def mlregistry(self):
         # use dummy models
-        from gistillery.ml import MlRegistry, RawTextProcessor, Summarizer, Tagger
+        from gistillery.ml import Summarizer, Tagger
+        from gistillery.preprocessing import RawTextProcessor
+        from gistillery.registry import MlRegistry
 
         class DummySummarizer(Summarizer):
             """Returns the first 10 characters of the input"""