podcast script generation component

.github/.devcontainer.json

@@ -7,5 +7,5 @@
       },
       "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"
 }

demo/app.py

@@ -1,9 +1,21 @@
 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.text_to_podcast import load_model
+from opennotebookllm.text_to_podcast import text_to_podcast
 
+PODCAST_PROMPT = """
+Convert this text into a podcast script.
+The conversation should be between 2 speakers.
+Use [SPEAKER1] and [SPEAKER2] to limit sections.
+Do not include [INTRO], [OUTRO] or any other [SECTION].
+Text:
+"""
+
+REPO = "allenai/OLMoE-1B-7B-0924-Instruct-GGUF"
 
 uploaded_file = st.file_uploader(
     "Choose a file", type=["pdf", "html", "txt", "docx", "md"]
@@ -17,9 +29,44 @@
     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]} . . .")
+
+    # I set this value as a quick safeguard but we should actually tokenize the text and count the number of real tokens.
+    if len(clean_text) > 4096 * 3:
+        st.warning(
+            f"Input text is too big ({len(clean_text)}). Using only a subset of it ({4096 * 3})."
+        )
+        clean_text = clean_text[: 4096 * 3]
+
+    model_name = st.selectbox(
+        "Select Model",
+        [
+            x
+            for x in list_repo_files(REPO)
+            if ".gguf" in x
+            # The float16 is too big for the 16GB RAM codespace
+            and "f16" not in x
+        ],
+        index=None,
+    )
+    if model_name:
+        with st.spinner("Downloading and Loading Model..."):
+            model = load_model(model_id=f"{REPO}/{model_name}")
+
+        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_podcast(
+                    clean_text, model, system_prompt=system_prompt.strip(), stream=True
+                ):
+                    text += chunk
+                    if text.endswith("\n"):
+                        st.write(text)
+                        text = ""

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"

src/opennotebookllm/text_to_podcast/__init__.py

@@ -0,0 +1,2 @@
+from .inference import load_model as load_model
+from .inference import text_to_podcast as text_to_podcast

src/opennotebookllm/text_to_podcast/inference.py

@@ -0,0 +1,32 @@
+from llama_cpp import Llama
+
+
+def load_model(
+    model_id: str = "allenai/OLMoE-1B-7B-0924-Instruct-GGUF/olmoe-1b-7b-0924-instruct-q8_0.gguf",
+) -> Llama:
+    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
+
+
+def text_to_podcast(
+    input_text: str, model: Llama, system_prompt: str, stream: bool = False
+):
+    response = model.create_chat_completion(
+        messages=[
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": input_text},
+        ],
+        stream=stream,
+    )
+    if stream:
+        for item in response:
+            if item["choices"][0].get("delta", {}).get("content", None):
+                yield item["choices"][0].get("delta", {}).get("content", None)
+    else:
+        return response["choices"][0]["message"]["content"]