tumblrbot 1.9.6__py3-none-any.whl → 1.10.0__py3-none-any.whl

tumblrbot/__main__.py

@@ -1,43 +1,47 @@
-from sys import exit as sys_exit
-
-from openai import OpenAI
-from rich.prompt import Confirm
-from rich.traceback import install
-
-from tumblrbot.flow.download import PostDownloader
-from tumblrbot.flow.examples import ExamplesWriter
-from tumblrbot.flow.fine_tune import FineTuner
-from tumblrbot.flow.generate import DraftGenerator
-from tumblrbot.utils.common import FlowClass
-from tumblrbot.utils.models import Tokens
-from tumblrbot.utils.tumblr import TumblrSession
-
-
-def main() -> None:
-    install()
-
-    tokens = Tokens.load()
-    with OpenAI(api_key=tokens.openai_api_key) as openai, TumblrSession(tokens) as tumblr:
-        if Confirm.ask("Download latest posts?", default=False):
-            PostDownloader(openai=openai, tumblr=tumblr).main()
-
-        examples_writer = ExamplesWriter(openai=openai, tumblr=tumblr)
-        if Confirm.ask("Create training data?", default=False):
-            examples_writer.main()
-
-        if Confirm.ask("Remove training data flagged by the OpenAI moderation? [bold]This can sometimes resolve errors with fine-tuning validation, but is slow.", default=False):
-            examples_writer.filter_examples()
-
-        fine_tuner = FineTuner(openai=openai, tumblr=tumblr)
-        fine_tuner.print_estimates()
-
-        message = "Resume monitoring the previous fine-tuning process?" if FlowClass.config.job_id else "Upload data to OpenAI for fine-tuning?"
-        if Confirm.ask(f"{message} [bold]You must do this to set the model to generate drafts from. Alternatively, manually enter a model into the config", default=False):
-            fine_tuner.main()
-
-        if Confirm.ask("Generate drafts?", default=False):
-            DraftGenerator(openai=openai, tumblr=tumblr).main()
-
-
-if __name__ == "__main__":
-    sys_exit(main())
+from locale import LC_ALL, setlocale
+from sys import exit as sys_exit
+from sys import maxsize
+
+from openai import OpenAI
+from rich.prompt import Confirm
+from rich.traceback import install
+
+from tumblrbot.flow.download import PostDownloader
+from tumblrbot.flow.examples import ExamplesWriter
+from tumblrbot.flow.fine_tune import FineTuner
+from tumblrbot.flow.generate import DraftGenerator
+from tumblrbot.utils.common import FlowClass
+from tumblrbot.utils.models import Tokens
+from tumblrbot.utils.tumblr import TumblrSession
+
+
+def main() -> None:
+    setlocale(LC_ALL, "")
+
+    install()
+
+    tokens = Tokens.load()
+    with OpenAI(api_key=tokens.openai_api_key, max_retries=maxsize) as openai, TumblrSession(tokens) as tumblr:
+        if Confirm.ask("Download latest posts?", default=False):
+            PostDownloader(openai, tumblr).main()
+
+        examples_writer = ExamplesWriter(openai, tumblr)
+        if Confirm.ask("Create training data?", default=False):
+            examples_writer.main()
+
+        if Confirm.ask("Remove training data flagged by the OpenAI moderation? [bold]This can sometimes resolve errors with fine-tuning validation, but is slow.", default=False):
+            examples_writer.filter_examples()
+
+        fine_tuner = FineTuner(openai, tumblr)
+        fine_tuner.print_estimates()
+
+        message = "Resume monitoring the previous fine-tuning process?" if FlowClass.config.job_id else "Upload data to OpenAI for fine-tuning?"
+        if Confirm.ask(f"{message} [bold]You must do this to set the model to generate drafts from. Alternatively, manually enter a model into the config", default=False):
+            fine_tuner.main()
+
+        if Confirm.ask("Generate drafts?", default=False):
+            DraftGenerator(openai, tumblr).main()
+
+
+if __name__ == "__main__":
+    sys_exit(main())

tumblrbot/flow/download.py

@@ -1,55 +1,55 @@
-from json import dump
-from typing import TYPE_CHECKING, override
-
-from tumblrbot.utils.common import FlowClass, PreviewLive
-from tumblrbot.utils.models import Post
-
-if TYPE_CHECKING:
-    from io import TextIOBase
-
-
-class PostDownloader(FlowClass):
-    @override
-    def main(self) -> None:
-        self.config.data_directory.mkdir(parents=True, exist_ok=True)
-
-        with PreviewLive() as live:
-            for blog_identifier in self.config.download_blog_identifiers:
-                data_path = self.get_data_path(blog_identifier)
-
-                completed = 0
-                after = 0
-                if data_path.exists():
-                    lines = data_path.read_bytes().splitlines() if data_path.exists() else []
-                    completed = len(lines)
-                    if lines:
-                        after = Post.model_validate_json(lines[-1]).timestamp
-
-                with data_path.open("a", encoding="utf_8") as fp:
-                    self.paginate_posts(
-                        blog_identifier,
-                        completed,
-                        after,
-                        fp,
-                        live,
-                    )
-
-    def paginate_posts(self, blog_identifier: str, completed: int, after: int, fp: TextIOBase, live: PreviewLive) -> None:
-        task_id = live.progress.add_task(f"Downloading posts from '{blog_identifier}'...", total=None, completed=completed)
-
-        while True:
-            response = self.tumblr.retrieve_published_posts(blog_identifier, after=after)
-            live.progress.update(task_id, total=response.response.blog.posts, completed=completed)
-
-            if not response.response.posts:
-                return
-
-            for post in response.response.posts:
-                dump(post, fp)
-                fp.write("\n")
-
-                model = Post.model_validate(post)
-                after = model.timestamp
-                live.custom_update(model)
-
-            completed += len(response.response.posts)
+from json import dump
+from typing import TYPE_CHECKING, override
+
+from tumblrbot.utils.common import FlowClass, PreviewLive
+from tumblrbot.utils.models import Post
+
+if TYPE_CHECKING:
+    from io import TextIOBase
+
+
+class PostDownloader(FlowClass):
+    @override
+    def main(self) -> None:
+        self.config.data_directory.mkdir(parents=True, exist_ok=True)
+
+        with PreviewLive() as live:
+            for blog_identifier in self.config.download_blog_identifiers:
+                data_path = self.get_data_path(blog_identifier)
+
+                completed = 0
+                after = 0
+                if data_path.exists():
+                    lines = data_path.read_bytes().splitlines() if data_path.exists() else []
+                    completed = len(lines)
+                    if lines:
+                        after = Post.model_validate_json(lines[-1]).timestamp
+
+                with data_path.open("a", encoding="utf_8") as fp:
+                    self.paginate_posts(
+                        blog_identifier,
+                        completed,
+                        after,
+                        fp,
+                        live,
+                    )
+
+    def paginate_posts(self, blog_identifier: str, completed: int, after: int, fp: TextIOBase, live: PreviewLive) -> None:
+        task_id = live.progress.add_task(f"Downloading posts from '{blog_identifier}'...", total=None, completed=completed)
+
+        while True:
+            response = self.tumblr.retrieve_published_posts(blog_identifier, after=after)
+            live.progress.update(task_id, total=response.response.blog.posts, completed=completed)
+
+            if not response.response.posts:
+                return
+
+            for post in response.response.posts:
+                dump(post, fp)
+                fp.write("\n")
+
+                model = Post.model_validate(post)
+                after = model.timestamp
+                live.custom_update(model)
+
+            completed += len(response.response.posts)

tumblrbot/flow/examples.py

@@ -1,97 +1,94 @@
-from collections.abc import Generator
-from itertools import batched
-from json import loads
-from math import ceil
-from re import IGNORECASE
-from re import compile as re_compile
-from typing import TYPE_CHECKING, override
-
-from openai import RateLimitError
-from rich import print as rich_print
-from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_random_exponential
-
-from tumblrbot.utils.common import FlowClass, PreviewLive
-from tumblrbot.utils.models import Example, Message, Post
-
-if TYPE_CHECKING:
-    from collections.abc import Generator, Iterable
-    from pathlib import Path
-
-    from openai._types import SequenceNotStr
-    from openai.types import ModerationCreateResponse, ModerationMultiModalInputParam
-
-
-class ExamplesWriter(FlowClass):
-    @override
-    def main(self) -> None:
-        self.config.examples_file.parent.mkdir(parents=True, exist_ok=True)
-
-        examples = [self.create_example(*prompt) for prompt in self.get_custom_prompts()]
-        examples.extend(self.create_example(self.config.user_message, str(post)) for post in self.get_valid_posts())
-        self.write_examples(examples)
-
-        rich_print(f"[bold]The examples file can be found at: '{self.config.examples_file}'\n")
-
-    def create_example(self, user_message: str, assistant_message: str) -> Example:
-        return Example(
-            messages=[
-                Message(role="developer", content=self.config.developer_message),
-                Message(role="user", content=user_message),
-                Message(role="assistant", content=assistant_message),
-            ],
-        )
-
-    def get_custom_prompts(self) -> Generator[tuple[str, str]]:
-        self.config.custom_prompts_file.parent.mkdir(parents=True, exist_ok=True)
-        self.config.custom_prompts_file.touch(exist_ok=True)
-
-        with self.config.custom_prompts_file.open("rb") as fp:
-            for line in fp:
-                data: dict[str, str] = loads(line)
-                yield from data.items()
-
-    # This function mostly exists to make writing examples atomic.
-    def write_examples(self, examples: Iterable[Example]) -> None:
-        with self.config.examples_file.open("w", encoding="utf_8") as fp:
-            for example in examples:
-                fp.write(f"{example.model_dump_json()}\n")
-
-    def get_valid_posts(self) -> Generator[Post]:
-        for path in self.get_data_paths():
-            posts = list(self.get_valid_posts_from_path(path))
-            yield from posts[-self.config.post_limit :]
-
-    def get_valid_posts_from_path(self, path: Path) -> Generator[Post]:
-        pattern = re_compile("|".join(self.config.filtered_words), IGNORECASE)
-        with path.open("rb") as fp:
-            for line in fp:
-                post = Post.model_validate_json(line)
-                if post.valid_text_post() and not (post.trail and self.config.filtered_words and pattern.search(str(post))):
-                    yield post
-
-    def filter_examples(self) -> None:
-        raw_examples = self.config.examples_file.read_bytes().splitlines()
-        old_examples = map(Example.model_validate_json, raw_examples)
-        new_examples: list[Example] = []
-        with PreviewLive() as live:
-            for batch in live.progress.track(
-                batched(old_examples, self.config.moderation_batch_size, strict=False),
-                ceil(len(raw_examples) / self.config.moderation_batch_size),
-                description="Removing flagged posts...",
-            ):
-                response = self.create_moderation_batch(tuple(map(Example.get_assistant_message, batch)))
-                new_examples.extend(example for example, moderation in zip(batch, response.results, strict=True) if not moderation.flagged)
-
-        self.write_examples(new_examples)
-
-        rich_print(f"[red]Removed {len(raw_examples) - len(new_examples)} posts.\n")
-
-    @retry(
-        stop=stop_after_attempt(10),
-        wait=wait_random_exponential(),
-        retry=retry_if_exception_type(RateLimitError),
-        before_sleep=lambda state: rich_print(f"[yellow]OpenAI rate limit exceeded. Waiting for {state.idle_for} seconds..."),
-        reraise=True,
-    )
-    def create_moderation_batch(self, api_input: str | SequenceNotStr[str] | Iterable[ModerationMultiModalInputParam]) -> ModerationCreateResponse:
-        return self.openai.moderations.create(input=api_input)
+from collections.abc import Generator
+from itertools import batched
+from json import loads
+from math import ceil
+from re import IGNORECASE
+from re import compile as re_compile
+from typing import TYPE_CHECKING, override
+
+from rich import print as rich_print
+from rich.console import Console
+
+from tumblrbot.utils.common import FlowClass, PreviewLive, localize_number
+from tumblrbot.utils.models import Example, Message, Post
+
+if TYPE_CHECKING:
+    from collections.abc import Generator, Iterable
+    from pathlib import Path
+
+    from openai._types import SequenceNotStr
+    from openai.types import ModerationCreateResponse, ModerationMultiModalInputParam
+
+
+class ExamplesWriter(FlowClass):
+    @override
+    def main(self) -> None:
+        self.config.examples_file.parent.mkdir(parents=True, exist_ok=True)
+
+        examples = [self.create_example(*prompt) for prompt in self.get_custom_prompts()]
+        examples.extend(self.create_example(self.config.user_message, str(post)) for post in self.get_valid_posts())
+        self.write_examples(examples)
+
+        rich_print(f"[bold]The examples file can be found at: '{self.config.examples_file}'\n")
+
+    def create_example(self, user_message: str, assistant_message: str) -> Example:
+        return Example(
+            messages=[
+                Message(role="developer", content=self.config.developer_message),
+                Message(role="user", content=user_message),
+                Message(role="assistant", content=assistant_message),
+            ],
+        )
+
+    def get_custom_prompts(self) -> Generator[tuple[str, str]]:
+        self.config.custom_prompts_file.parent.mkdir(parents=True, exist_ok=True)
+        self.config.custom_prompts_file.touch(exist_ok=True)
+
+        with self.config.custom_prompts_file.open("rb") as fp:
+            for line in fp:
+                data: dict[str, str] = loads(line)
+                yield from data.items()
+
+    # This function mostly exists to make writing examples (mostly) atomic.
+    # If there is an error dumping the models or writing to the file,
+    # then it will leave a partially written or empty file behind.
+    def write_examples(self, examples: Iterable[Example]) -> None:
+        with self.config.examples_file.open("w", encoding="utf_8") as fp:
+            for example in examples:
+                fp.write(f"{example.model_dump_json()}\n")
+
+    def get_valid_posts(self) -> Generator[Post]:
+        for path in self.get_data_paths():
+            if path.exists():
+                posts = list(self.get_valid_posts_from_path(path))
+                yield from posts[-self.config.post_limit :]
+            else:
+                Console(stderr=True, style="logging.level.warning").print(f"{path} does not exist!")
+
+    def get_valid_posts_from_path(self, path: Path) -> Generator[Post]:
+        pattern = re_compile("|".join(self.config.filtered_words), IGNORECASE)
+        with path.open("rb") as fp:
+            for line in fp:
+                post = Post.model_validate_json(line)
+                if post.valid_text_post() and not (post.trail and self.config.filtered_words and pattern.search(str(post))):
+                    yield post
+
+    def filter_examples(self) -> None:
+        raw_examples = self.config.examples_file.read_bytes().splitlines()
+        old_examples = map(Example.model_validate_json, raw_examples)
+        new_examples: list[Example] = []
+        with PreviewLive() as live:
+            for batch in live.progress.track(
+                batched(old_examples, self.config.moderation_batch_size, strict=False),
+                ceil(len(raw_examples) / self.config.moderation_batch_size),
+                description="Removing flagged posts...",
+            ):
+                response = self.create_moderation_batch(tuple(map(Example.get_assistant_message, batch)))
+                new_examples.extend(example for example, moderation in zip(batch, response.results, strict=True) if not moderation.flagged)
+
+        self.write_examples(new_examples)
+
+        rich_print(f"[red]Removed {localize_number(len(raw_examples) - len(new_examples))} posts.\n")
+
+    def create_moderation_batch(self, api_input: str | SequenceNotStr[str] | Iterable[ModerationMultiModalInputParam]) -> ModerationCreateResponse:
+        return self.openai.moderations.create(input=api_input)