From 3fc7ff1d39329420e94c082346647196c0794dc3 Mon Sep 17 00:00:00 2001 From: David de la Iglesia Castro Date: Mon, 25 Nov 2024 13:54:48 +0100 Subject: [PATCH] podcast script generation component (#15) * Add devcontainer and requirements * Add pyproject.toml * Add data_loaders and tests * Add data_cleaners and tests * Update demo * Add `LOADERS` and `CLEANERS` * Add markdown and docx * Add API Reference * Update tests * Update install * Add initial scripts * More tests * fix merge * Add podcast writing to demo/app * Add missing deps * Add text_to_podcast module * Expose model options and prompt tuning in the app * pre-commit * Strip system_prompt * Rename to inference module. Add docstrings * pre-commit * Add CURATED_REPOS * JSON prompt * Update API docs * Fix format * Make text cutoff based on `model.n_ctx()`. Consider ~4 characters per token as a resonable default. * Add inference tests * Drop __init__ imports * Fix outdated arg * Drop redundant JSON output in prompt * Update default stop --- .github/.devcontainer.json | 5 +- demo/app.py | 63 +++++++++++++- docs/api.md | 4 + pyproject.toml | 9 +- src/opennotebookllm/inference/__init__.py | 0 .../inference/model_loaders.py | 28 +++++++ src/opennotebookllm/inference/text_to_text.py | 84 +++++++++++++++++++ ...d_clean.py => test_data_load_and_clean.py} | 0 .../test_model_load_and_inference.py | 68 +++++++++++++++ tests/unit/inference/test_model_loaders.py | 12 +++ .../preprocessing/test_data_cleaners.py | 0 .../preprocessing/test_data_loaders.py | 0 12 files changed, 263 insertions(+), 10 deletions(-) create mode 100644 src/opennotebookllm/inference/__init__.py create mode 100644 src/opennotebookllm/inference/model_loaders.py create mode 100644 src/opennotebookllm/inference/text_to_text.py rename tests/integration/{test_load_and_clean.py => test_data_load_and_clean.py} (100%) create mode 100644 tests/integration/test_model_load_and_inference.py create mode 100644 tests/unit/inference/test_model_loaders.py rename tests/{ => unit}/preprocessing/test_data_cleaners.py (100%) rename tests/{ => unit}/preprocessing/test_data_loaders.py (100%) diff --git a/.github/.devcontainer.json b/.github/.devcontainer.json index b467d3b..f26fe37 100644 --- a/.github/.devcontainer.json +++ b/.github/.devcontainer.json @@ -4,8 +4,7 @@ "features": { "ghcr.io/devcontainers/features/python": { "version": "latest" - }, - "packages": ["libgl1-mesa-dev"] + } }, - "postCreateCommand": "pip install -e '.[demo]'" + "postCreateCommand": "pip install -e . --extra-index-url https://abetlen.github.io/llama-cpp-python/whl/cpu" } diff --git a/demo/app.py b/demo/app.py index e4046f3..ee1ad3f 100644 --- a/demo/app.py +++ b/demo/app.py @@ -1,9 +1,32 @@ from pathlib import Path import streamlit as st +from huggingface_hub import list_repo_files from opennotebookllm.preprocessing import DATA_LOADERS, DATA_CLEANERS +from opennotebookllm.inference.model_loaders import load_llama_cpp_model +from opennotebookllm.inference.text_to_text import text_to_text_stream +PODCAST_PROMPT = """ +You are a helpful podcast writer. +You will take the input text and generate a conversation between 2 speakers. +Example of response: +{ + "Speaker 1": "Welcome to our podcast, where we explore the latest advancements in AI and technology. I'm your host, and today we're going to dive into the exciting world of TrustWorthy AI.", + "Speaker 2": "Hi, I'm excited to be here, so what is TrustWorthy AI?", + "Speaker 1":"Ah, great question! It is a term used by the European High Level Expert Group on AI. Mozilla defines trustworthy AI as AI that is demonstrably worthy of trust, tech that considers accountability, agency, and individual and collective well-being." +} +""" + +CURATED_REPOS = [ + "allenai/OLMoE-1B-7B-0924-Instruct-GGUF", + "MaziyarPanahi/SmolLM2-1.7B-Instruct-GGUF", + # system prompt seems to be ignored for this model. + # "microsoft/Phi-3-mini-4k-instruct-gguf", + "HuggingFaceTB/SmolLM2-360M-Instruct-GGUF", + "Qwen/Qwen2.5-1.5B-Instruct-GGUF", + "Qwen/Qwen2.5-3B-Instruct-GGUF", +] uploaded_file = st.file_uploader( "Choose a file", type=["pdf", "html", "txt", "docx", "md"] @@ -17,9 +40,45 @@ raw_text = DATA_LOADERS[extension](uploaded_file) with col1: st.title("Raw Text") - st.write(raw_text[:200]) + st.text_area(f"Total Length: {len(raw_text)}", f"{raw_text[:500]} . . .") clean_text = DATA_CLEANERS[extension](raw_text) with col2: st.title("Cleaned Text") - st.write(clean_text[:200]) + st.text_area(f"Total Length: {len(clean_text)}", f"{clean_text[:500]} . . .") + + repo_name = st.selectbox("Select Repo", CURATED_REPOS) + model_name = st.selectbox( + "Select Model", + [ + x + for x in list_repo_files(repo_name) + if ".gguf" in x.lower() and ("q8" in x.lower() or "fp16" in x.lower()) + ], + index=None, + ) + if model_name: + with st.spinner("Downloading and Loading Model..."): + model = load_llama_cpp_model(model_id=f"{repo_name}/{model_name}") + + # ~4 characters per token is considered a reasonable default. + max_characters = model.n_ctx() * 4 + if len(clean_text) > max_characters: + st.warning( + f"Input text is too big ({len(clean_text)})." + f" Using only a subset of it ({max_characters})." + ) + clean_text = clean_text[:max_characters] + + system_prompt = st.text_area("Podcast generation prompt", value=PODCAST_PROMPT) + + if st.button("Generate Podcast Script"): + with st.spinner("Generating Podcast Script..."): + text = "" + for chunk in text_to_text_stream( + clean_text, model, system_prompt=system_prompt.strip() + ): + text += chunk + if text.endswith("\n"): + st.write(text) + text = "" diff --git a/docs/api.md b/docs/api.md index c304571..a7c1c65 100644 --- a/docs/api.md +++ b/docs/api.md @@ -1,3 +1,7 @@ # API Reference ::: opennotebookllm.preprocessing.data_cleaners + +::: opennotebookllm.inference.model_loaders + +::: opennotebookllm.inference.text_to_text diff --git a/pyproject.toml b/pyproject.toml index 704bf32..ac35def 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,9 +10,12 @@ requires-python = ">=3.10" dynamic = ["version"] dependencies = [ "beautifulsoup4", + "huggingface-hub", + "llama-cpp-python", "loguru", "PyPDF2[crypto]", - "python-docx" + "python-docx", + "streamlit", ] [project.optional-dependencies] @@ -27,10 +30,6 @@ tests = [ "pytest-sugar>=0.9.6", ] -demo = [ - "streamlit" -] - [project.urls] Documentation = "https://mozilla-ai.github.io/OpenNotebookLLM/" Issues = "https://github.com/mozilla-ai/OpenNotebookLLM/issues" diff --git a/src/opennotebookllm/inference/__init__.py b/src/opennotebookllm/inference/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/opennotebookllm/inference/model_loaders.py b/src/opennotebookllm/inference/model_loaders.py new file mode 100644 index 0000000..f313d45 --- /dev/null +++ b/src/opennotebookllm/inference/model_loaders.py @@ -0,0 +1,28 @@ +from llama_cpp import Llama + + +def load_llama_cpp_model( + model_id: str, +) -> Llama: + """ + Loads the given model_id using Llama.from_pretrained. + + Examples: + >>> model = load_model( + "allenai/OLMoE-1B-7B-0924-Instruct-GGUF/olmoe-1b-7b-0924-instruct-q8_0.gguf") + + Args: + model_id (str): The model id to load. + Format is expected to be `{org}/{repo}/{filename}`. + + Returns: + Llama: The loaded model. + """ + org, repo, filename = model_id.split("/") + model = Llama.from_pretrained( + repo_id=f"{org}/{repo}", + filename=filename, + # 0 means that the model limit will be used, instead of the default (512) or other hardcoded value + n_ctx=0, + ) + return model diff --git a/src/opennotebookllm/inference/text_to_text.py b/src/opennotebookllm/inference/text_to_text.py new file mode 100644 index 0000000..6f2175a --- /dev/null +++ b/src/opennotebookllm/inference/text_to_text.py @@ -0,0 +1,84 @@ +from typing import Iterator + +from llama_cpp import Llama + + +def chat_completion( + input_text: str, + model: Llama, + system_prompt: str, + return_json: bool, + stream: bool, + stop: str | list[str] | None = None, +) -> str | Iterator[str]: + # create_chat_completion uses an empty list as default + stop = stop or [] + return model.create_chat_completion( + messages=[ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": input_text}, + ], + response_format={ + "type": "json_object", + } + if return_json + else None, + stream=stream, + stop=stop, + ) + + +def text_to_text( + input_text: str, + model: Llama, + system_prompt: str, + return_json: bool = True, + stop: str | list[str] | None = None, +) -> str: + """ + Transforms input_text using the given model and system prompt. + + Args: + input_text (str): The text to be transformed. + model (Llama): The model to use for conversion. + system_prompt (str): The system prompt to use for conversion. + return_json (bool, optional): Whether to return the response as JSON. + Defaults to True. + stop (str | list[str] | None, optional): The stop token(s). + + Returns: + str: The full transformed text. + """ + response = chat_completion( + input_text, model, system_prompt, return_json, stop=stop, stream=False + ) + return response["choices"][0]["message"]["content"] + + +def text_to_text_stream( + input_text: str, + model: Llama, + system_prompt: str, + return_json: bool = True, + stop: str | list[str] | None = None, +) -> Iterator[str]: + """ + Transforms input_text using the given model and system prompt. + + Args: + input_text (str): The text to be transformed. + model (Llama): The model to use for conversion. + system_prompt (str): The system prompt to use for conversion. + return_json (bool, optional): Whether to return the response as JSON. + Defaults to True. + stop (str | list[str] | None, optional): The stop token(s). + + Yields: + str: Chunks of the transformed text as they are available. + """ + response = chat_completion( + input_text, model, system_prompt, return_json, stop=stop, stream=True + ) + for item in response: + if item["choices"][0].get("delta", {}).get("content", None): + yield item["choices"][0].get("delta", {}).get("content", None) diff --git a/tests/integration/test_load_and_clean.py b/tests/integration/test_data_load_and_clean.py similarity index 100% rename from tests/integration/test_load_and_clean.py rename to tests/integration/test_data_load_and_clean.py diff --git a/tests/integration/test_model_load_and_inference.py b/tests/integration/test_model_load_and_inference.py new file mode 100644 index 0000000..7a248e7 --- /dev/null +++ b/tests/integration/test_model_load_and_inference.py @@ -0,0 +1,68 @@ +import json +from typing import Iterator + +import pytest + +from opennotebookllm.inference.model_loaders import load_llama_cpp_model +from opennotebookllm.inference.text_to_text import text_to_text, text_to_text_stream + + +def test_model_load_and_inference_text_to_text(): + model = load_llama_cpp_model( + "HuggingFaceTB/smollm-135M-instruct-v0.2-Q8_0-GGUF/smollm-135m-instruct-add-basics-q8_0.gguf" + ) + result = text_to_text( + "What is the capital of France?", + model=model, + system_prompt="", + ) + assert isinstance(result, str) + assert json.loads(result)["Capital"] == "Paris" + + +def test_model_load_and_inference_text_to_text_no_json(): + model = load_llama_cpp_model( + "HuggingFaceTB/smollm-135M-instruct-v0.2-Q8_0-GGUF/smollm-135m-instruct-add-basics-q8_0.gguf" + ) + result = text_to_text( + "What is the capital of France?", + model=model, + system_prompt="", + return_json=False, + stop=".", + ) + assert isinstance(result, str) + with pytest.raises(json.JSONDecodeError): + json.loads(result) + assert result.startswith("The capital of France is Paris") + + +def test_model_load_and_inference_text_to_text_stream(): + model = load_llama_cpp_model( + "HuggingFaceTB/smollm-135M-instruct-v0.2-Q8_0-GGUF/smollm-135m-instruct-add-basics-q8_0.gguf" + ) + result = text_to_text_stream( + "What is the capital of France?", + model=model, + system_prompt="", + ) + assert isinstance(result, Iterator) + assert json.loads("".join(result))["Capital"] == "Paris" + + +def test_model_load_and_inference_text_to_text_stream_no_json(): + model = load_llama_cpp_model( + "HuggingFaceTB/smollm-135M-instruct-v0.2-Q8_0-GGUF/smollm-135m-instruct-add-basics-q8_0.gguf" + ) + result = text_to_text_stream( + "What is the capital of France?", + model=model, + system_prompt="", + return_json=False, + stop=".", + ) + assert isinstance(result, Iterator) + result = "".join(result) + with pytest.raises(json.JSONDecodeError): + json.loads(result) + assert result.startswith("The capital of France is Paris") diff --git a/tests/unit/inference/test_model_loaders.py b/tests/unit/inference/test_model_loaders.py new file mode 100644 index 0000000..630c62a --- /dev/null +++ b/tests/unit/inference/test_model_loaders.py @@ -0,0 +1,12 @@ +from llama_cpp import Llama + +from opennotebookllm.inference.model_loaders import load_llama_cpp_model + + +def test_load_llama_cpp_model(): + model = load_llama_cpp_model( + "HuggingFaceTB/smollm-135M-instruct-v0.2-Q8_0-GGUF/smollm-135m-instruct-add-basics-q8_0.gguf" + ) + assert isinstance(model, Llama) + # we set n_ctx=0 to indicate that we want to use the model's default context + assert model.n_ctx() == 2048 diff --git a/tests/preprocessing/test_data_cleaners.py b/tests/unit/preprocessing/test_data_cleaners.py similarity index 100% rename from tests/preprocessing/test_data_cleaners.py rename to tests/unit/preprocessing/test_data_cleaners.py diff --git a/tests/preprocessing/test_data_loaders.py b/tests/unit/preprocessing/test_data_loaders.py similarity index 100% rename from tests/preprocessing/test_data_loaders.py rename to tests/unit/preprocessing/test_data_loaders.py