Spaces:

jugarte00
/

wsws-chatbot

Sleeping

+#!/bin/bash
+# Usage: svc service [command]
+# Example: svc app ps aux
+# Validate arguments
+if [[ $# -lt 1 ]]; then
+    echo "Usage: svc service [command]"
+    exit 1
+fi
+# Service
+case $1 in
+    *)
+        SVC=
+        CMD=
+        ;;
+esac
+CMD_ARGS="${@:2}"
+if [[ -z "$CMD_ARGS" ]]; then
+    CMD_ARGS="$CMD"
+fi
+docker compose exec $SVC $CMD_ARGS

.devcontainer/devcontainer.json ADDED Viewed

	@@ -0,0 +1,51 @@

+// For format details, see https://aka.ms/devcontainer.json. For config options, see the
+// README at: https://github.com/devcontainers/templates/tree/main/src/python
+{
+    "name": "Article Embedding",
+    "dockerComposeFile": "docker-compose.yml",
+    "service": "dev",
+    "runServices": [
+        "dev",
+        "qdrant"
+    ],
+    "workspaceFolder": "/workspaces/article-embedding",
+    "features": {
+        "ghcr.io/devcontainers/features/github-cli:1": {},
+        "ghcr.io/nikobockerman/devcontainer-features/poetry-persistent-cache:1": {},
+        "ghcr.io/devcontainers-extra/features/poetry:2": {},
+        "ghcr.io/devcontainers/features/docker-outside-of-docker:1": {}
+    },
+    "containerEnv": {
+        "DOCKER_CLI_HINTS": "false",
+        "POETRY_VIRTUALENVS_IN_PROJECT": "false",
+        "VIRTUAL_ENV": "/mnt/poetry-persistent-cache/virtualenvs/article-embedding-p7H3l83p-py3.12",
+        "COMPOSE_PROJECT_NAME": "article-embedding_devcontainer",
+        "COMPOSE_FILE": "${containerWorkspaceFolder}/.devcontainer/docker-compose.yml"
+    },
+    "remoteEnv": {
+        "PATH": "${containerEnv:VIRTUAL_ENV}/bin:${containerWorkspaceFolder}/.devcontainer/dev/bin:${containerEnv:PATH}"
+    },
+    "mounts": [
+        "type=bind,source=${localEnv:HOME}${localEnv:USERPROFILE}/.gitconfig,target=/home/vscode/.gitconfig,readonly",
+        "type=bind,source=${localEnv:HOME}${localEnv:USERPROFILE}/.ssh,target=/home/vscode/.ssh,readonly"
+    ],
+    "postCreateCommand": ".devcontainer/scripts/post-create.sh ${containerWorkspaceFolder}",
+    "postStartCommand": ".devcontainer/scripts/post-start.sh ${containerWorkspaceFolder}",
+    "customizations": {
+        "vscode": {
+            "extensions": [
+                "ms-python.debugpy",
+                "github.vscode-github-actions",
+                "matangover.mypy",
+                "tamasfe.even-better-toml",
+                "humao.rest-client",
+                "charliermarsh.ruff",
+                "ms-toolsai.jupyter",
+                "redhat.vscode-yaml"
+            ],
+            "settings": {
+                "python.defaultInterpreterPath": "${containerEnv:VIRTUAL_ENV}/bin/python"
+            }
+        }
+    }
+}

.devcontainer/docker-compose.yml ADDED Viewed

	@@ -0,0 +1,19 @@

+services:
+  dev:
+    profiles:
+      - devcontainer
+    build: dev
+    volumes:
+      - ../..:/workspaces:cached
+      - ..:/workspaces/article-embedding:cached
+    command: sleep infinity
+  qdrant:
+    image: qdrant/qdrant
+    ports:
+      - "6333:6333"
+    volumes:
+      - qdrant:/qdrant/storage:z
+volumes:
+  qdrant:

.devcontainer/scripts/post-create.sh ADDED Viewed

	@@ -0,0 +1,25 @@

+#!/bin/bash
+# Fail on any error
+set -e
+# Ensure the lock file is up to date, otherwise the install step may fail later
+poetry lock --no-update
+# Remove the old-style .venv symlink if it exists
+if [[ -L .venv ]]; then
+  rm .venv
+fi
+# Force the virtual environment to the path set in the .devcontainer.json file
+python -mvenv $VIRTUAL_ENV
+PATH="$VIRTUAL_ENV/bin:$PATH"
+poetry install
+# Activate pre commit hooks
+pre-commit install
+# Configuration files
+cat >> ~/.inputrc <<EOF
+set completion-ignore-case on
+set editing-mode vi
+EOF

.devcontainer/scripts/post-start.sh ADDED Viewed

	@@ -0,0 +1,2 @@


1	+ #!/bin/bash
2	+ set -e

.gitattributes CHANGED Viewed

@@ -1,35 +1,5 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text

+poetry.lock binary
+*.ipynb filter=nbstripout
+*.zpln filter=nbstripout
+*.ipynb diff=ipynb

.github/workflows/update_space.yml ADDED Viewed

	@@ -0,0 +1,32 @@

+name: Run Python script
+on:
+  push:
+    branches:
+    - main
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+    - name: Checkout
+      uses: actions/checkout@v4
+    - name: Install Poetry
+      run: pipx install poetry
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: "3.11"
+        cache: "poetry"
+    - name: Deploy
+      run: |-
+        pipx inject poetry poetry-plugin-export
+        poetry export --only=main >requirements.txt
+        poetry install --only=ci --no-root -q
+        PATH=$(poetry env info --path)/bin:$PATH
+        huggingface-cli login --token ${{ secrets.HF_TOKEN }}
+        gradio deploy

.gitignore ADDED Viewed

	@@ -0,0 +1,7 @@

+.*checkpoint
+.keys/
+__pycache__/
+.gradio/
+.env
+data/
+.venv/

.pre-commit-config.yaml ADDED Viewed

	@@ -0,0 +1,19 @@

+# See https://pre-commit.com for more information
+# See https://pre-commit.com/hooks.html for more hooks
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v3.2.0
+    hooks:
+      - id: trailing-whitespace
+      - id: end-of-file-fixer
+      - id: check-yaml
+      - id: check-added-large-files
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    # Ruff version.
+    rev: v0.7.4
+    hooks:
+      # Run the linter.
+      - id: ruff
+        args: [--fix]
+      # Run the formatter.
+      - id: ruff-format

.vscode/launch.json ADDED Viewed

	@@ -0,0 +1,59 @@

+{
+    // Use IntelliSense to learn about possible attributes.
+    // Hover to view descriptions of existing attributes.
+    // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "App",
+            "type": "debugpy",
+            "request": "launch",
+            "module": "article_embedding",
+            "args": [
+                "app"
+            ],
+            "console": "integratedTerminal",
+            "justMyCode": false,
+        },
+        {
+            "name": "Scanner",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "article_embedding/couchdb.py",
+            "console": "integratedTerminal",
+            "justMyCode": false,
+        },
+        {
+            "name": "Loader",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "article_embedding/loader.py",
+            "console": "integratedTerminal",
+            "justMyCode": false,
+        },
+        {
+            "name": "Embedder",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "article_embedding/embed.py",
+            "console": "integratedTerminal",
+            "justMyCode": false,
+        },
+        {
+            "name": "Query",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "article_embedding/query.py",
+            "console": "integratedTerminal",
+            "justMyCode": false,
+        },
+        {
+            "name": "Benchmark",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "article_embedding/benchmark.py",
+            "console": "integratedTerminal",
+            "justMyCode": false,
+        },
+    ]
+}

.vscode/settings.json ADDED Viewed

	@@ -0,0 +1,26 @@

+{
+    "editor.codeActionsOnSave": {
+        "source.organizeImports": "always",
+        "source.unusedImports": "always"
+    },
+    "editor.formatOnPaste": true,
+    "editor.formatOnSave": true,
+    "editor.formatOnSaveMode": "file",
+    "editor.formatOnType": true,
+    "editor.renderWhitespace": "all",
+    "editor.rulers": [
+        132
+    ],
+    "files.trimTrailingWhitespace": true,
+    "mypy.runUsingActiveInterpreter": true,
+    "python.analysis.importFormat": "absolute",
+    "python.analysis.autoFormatStrings": true,
+    "python.analysis.autoImportCompletions": true,
+    "ruff.organizeImports": true,
+    "[jsonc]": {
+        "editor.defaultFormatter": "vscode.json-language-features",
+    },
+    "[python]": {
+        "editor.defaultFormatter": "charliermarsh.ruff",
+    }
+}

README.md CHANGED Viewed

@@ -1,13 +1,148 @@
 ---
-title: Wsws Chatbot
-emoji: 💻
-colorFrom: red
-colorTo: green
 sdk: gradio
 sdk_version: 5.12.0
-app_file: app.py
 pinned: false
 short_description: Ask the WSWS chatbot
 ---
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

 ---
+title: wsws-chatbot
+emoji: 📊
+colorFrom: green
+colorTo: indigo
 sdk: gradio
 sdk_version: 5.12.0
+app_file: article_embedding/app/ui.py
 pinned: false
 short_description: Ask the WSWS chatbot
 ---
+## Filter
+Identify the top 60 most important articles of the year
+| Variables            | Description                           |
+| -------------------- | ------------------------------------- |
+| articlesPerYear      | Number of articles per year           |
+| perspectivePercent   | Percentage of perspectives            |
+| usPercent            | Percentage of US articles             |
+| maxArticlesPerAuthor | Maximum number of articles per author |
+### Important topics
+| Topic ID                             | Topic Name               |
+| ------------------------------------ | ------------------------ |
+| 3f1db5b4-7915-4cef-a132-3b2e51f0728d | US Politics              |
+| 105b9ec2-0cdd-480d-9c55-6f806fa7b9bc | Perspectives             |
+| 04d2a008-8cef-4928-9a38-86b205a7590f | The coronavirus pandemic |
+| ea69254c-a976-4e67-a79c-843202a10e55 | United States            |
+| b0e983c5-6e85-4829-a9fe-3c25f02c211f | North America            |
+## Analysis
+1. Document metadata structure
+- The metadata is stored as JSON
+- Yes, there are structured attributes in the metadata
+- There are many other structured attributes but none that are relevant to the system design
+2. Vector database setup
+- Qdrant
+- Documents are embedded using the Stella 400M v5 embedding
+3. Integration of filters and constraints
+- Excellent question. The constraints need not be strict. They should be approximately met, since relevance to the provided queries is also important
+- The documents should be selected first based on relevance and then adjust the final selection to meet the constraints. It's ok to select less relevant articles to get closer to the constraints, but up to a certain point.
+4. Date range filtering
+- The date range is strict
+5. Result volume and ranking logic
+- Great question too. The articles should be ranked based on their relevance, with all queries adding documents to the result set, but choosing the most relevant ones from this pool. For example, it would be ok if all articles end up being derived from only one query.
+6. Combining multiple queries
+- The former. Each query is run separately and then the results are combined based on the document relevance score. Since the same document may be returned for different queries, in each case with a different relevance score, the highest relevance score should be used.
+7. User preferences and constraint handling
+- Exactly, the result set should be post processed to meet constraints
+8. Performance considerations
+- The document database has about 100,000 documents
+- No latency requirements. The accuracy of the system is the only requirement
+9. Evaluation and Benchmarks
+- The system success will be measured based on relevance judgments and coverage of constraints, in that order
+- A feedback loop to improve the retrieval strategies or embedding approach is definitely on the table, if needed.
+        self.embedding_model = SentenceTransformerModel("dunzhang/stella_en_400M_v5")
+# Performance benchmarks
+The script at `article_embedding/modal_app.py` was used to measure the performance of the embedding model. The results are as follows:
+|  GPU   |   Speed |  Cost |
+| :----: | ------: | ----: |
+|   T4   | 130.0ms | 2.13¢ |
+|   L4   |  82.7ms | 1.84¢ |
+| M3 Max |  60.0ms |   N/A |
+|  A10G  |  64.2ms | 1.96¢ |
+|  A100  |  42.2ms | 3.26¢ |
+|  H100  |  24.9ms | 3.15¢ |
+These are benchmarks on the M3 Max 128GB:
+| Model          |   f32 |   f16 | Context |  Size |
+| -------------- | ----: | ----: | ------: | ----: |
+| jasper sdpa    | 373ms | 314ms |   1,024 | 1,024 |
+| jasper         | 354ms | 330ms |   1,024 | 1,024 |
+| stella-en-400M |  65ms |  62ms | 131,072 | 1,024 |
+## Models
+[Jasper](https://huggingface.co/infgrad/jasper_en_vision_language_v1)
+## Issues
+```sh
+$  cd /Users/jlopez/git-repos/icfi/article-embedding ; /usr/bin/env /Users/jlopez/git-repos/icfi/article-embedding/.venv/bin/python /Users/jlopez/.vscode/extensions/ms-python.debugpy-2024.10.0-darwin-arm64/bundled/libs/debugpy/adapter/../../debugpy/launcher 63791 -- article_embedding/loader.py
+  0%|                                                                                                                                                                               | 0/314855 [00:00<?, ?it/s]E+00379.644: /handling disconnect from Adapter/
+             Failed to kill Debuggee[PID=93387]
+             Traceback (most recent call last):
+               File "/Users/jlopez/.vscode/extensions/ms-python.debugpy-2024.10.0-darwin-arm64/bundled/libs/debugpy/adapter/../../debugpy/launcher/../../debugpy/launcher/debuggee.py", line 190, in kill
+                 os.killpg(process.pid, signal.SIGKILL)
+             PermissionError: [Errno 1] Operation not permitted
+             Stack where logged:
+               File "/Users/jlopez/.pyenv/versions/3.12.5/lib/python3.12/threading.py", line 1032, in _bootstrap
+                 self._bootstrap_inner()
+               File "/Users/jlopez/.pyenv/versions/3.12.5/lib/python3.12/threading.py", line 1075, in _bootstrap_inner
+                 self.run()
+               File "/Users/jlopez/.pyenv/versions/3.12.5/lib/python3.12/threading.py", line 1012, in run
+                 self._target(*self._args, **self._kwargs)
+               File "/Users/jlopez/.vscode/extensions/ms-python.debugpy-2024.10.0-darwin-arm64/bundled/libs/debugpy/adapter/../../debugpy/launcher/../../debugpy/common/messaging.py", line 1458, in _run_handlers
+                 handler()
+               File "/Users/jlopez/.vscode/extensions/ms-python.debugpy-2024.10.0-darwin-arm64/bundled/libs/debugpy/adapter/../../debugpy/launcher/../../debugpy/common/messaging.py", line 1488, in _handle_disconnect
+                 handler()
+               File "/Users/jlopez/.vscode/extensions/ms-python.debugpy-2024.10.0-darwin-arm64/bundled/libs/debugpy/adapter/../../debugpy/launcher/../../debugpy/launcher/handlers.py", line 152, in disconnect
+                 debuggee.kill()
+               File "/Users/jlopez/.vscode/extensions/ms-python.debugpy-2024.10.0-darwin-arm64/bundled/libs/debugpy/adapter/../../debugpy/launcher/../../debugpy/launcher/debuggee.py", line 192, in kill
+                 log.swallow_exception("Failed to kill {0}", describe())
+               File "/Users/jlopez/.vscode/extensions/ms-python.debugpy-2024.10.0-darwin-arm64/bundled/libs/debugpy/adapter/../../debugpy/launcher/../../debugpy/common/log.py", line 215, in swallow_exception
+                 _exception(format_string, *args, **kwargs)
+E+00379.648: Failed to kill Debuggee[PID=93387]
+             Traceback (most recent call last):
+               File "/Users/jlopez/.vscode/extensions/ms-python.debugpy-2024.10.0-darwin-arm64/bundled/libs/debugpy/adapter/../../debugpy/launcher/../../debugpy/launcher/debuggee.py", line 190, in kill
+                 os.killpg(process.pid, signal.SIGKILL)
+             PermissionError: [Errno 1] Operation not permitted
+             Stack where logged:
+               File "/Users/jlopez/.vscode/extensions/ms-python.debugpy-2024.10.0-darwin-arm64/bundled/libs/debugpy/adapter/../../debugpy/launcher/../../debugpy/launcher/debuggee.py", line 192, in kill
+                 log.swallow_exception("Failed to kill {0}", describe())
+               File "/Users/jlopez/.vscode/extensions/ms-python.debugpy-2024.10.0-darwin-arm64/bundled/libs/debugpy/adapter/../../debugpy/launcher/../../debugpy/common/log.py", line 215, in swallow_exception
+                 _exception(format_string, *args, **kwargs)
+```

article_embedding/__init__.py ADDED Viewed

File without changes

article_embedding/__main__.py ADDED Viewed

	@@ -0,0 +1,4 @@

+from article_embedding.cli import cli
+if __name__ == "__main__":
+    cli()

article_embedding/app/__init__.py ADDED Viewed

	@@ -0,0 +1,3 @@


1	+ from .ui import run_app
2	+
3	+ __all__ = ["run_app"]

article_embedding/app/gather.py ADDED Viewed

	@@ -0,0 +1,71 @@

+from typing import Any
+import gradio as gr
+from article_embedding.retrieval import retrieve_documents
+DEFAULT_QUERY = """
+The development of the COVID-19 pandemic globally.
+The “herd immunity” and “mitigationist” response of capitalist governments globally to the COVID-19 pandemic.
+The only viable solution of global elimination of the COVID-19 pandemic.
+The science of the COVID-19 pandemic: Long COVID, airborne transmission, viral evolution.
+""".strip()
+def gather_articles(
+    query: str,
+    year: int,
+    articles_per_year: int,
+    pct_perspectives: int,
+    pct_us: int,
+    max_per_author: int,
+) -> list[list[Any]]:
+    queries = [q.strip() for q in query.split("\n") if q.strip()]
+    docs = retrieve_documents(
+        queries=queries,
+        start_date=f"{year}-01-01",
+        end_date=f"{year + 1}-01-01",
+        number_of_articles=articles_per_year,
+        op_ed_ratio=pct_perspectives / 100.0,
+        us_ratio=pct_us / 100.0,
+        max_per_author=max_per_author,
+    )
+    return [
+        [
+            ix + 1,
+            d.query_id,
+            f"{d.score * 100:.1f}%",
+            d.meta.published.strftime("%Y/%m/%d") if d.meta.published else "N/A",
+            ", ".join(d.meta.authors)[:32] if d.meta.authors else "N/A",
+            f'[{d.meta.title or "N/A"}]({d.meta.url})',
+        ]
+        for ix, d in enumerate(docs)
+    ]
+def gather_tab() -> None:
+    with gr.Tab("Gather"):
+        with gr.Row():
+            with gr.Column(scale=2):
+                query = gr.Textbox(label="Query", placeholder="Enter multiple queries, one per line", value=DEFAULT_QUERY)
+                gather = gr.Button(value="Gather articles")
+            with gr.Column():
+                year = gr.Slider(minimum=2021, maximum=2025, value=2021, label="Year", step=1)
+                articles_per_year = gr.Slider(label="Articles per year", minimum=1, maximum=100, value=60, step=1)
+                pct_perspectives = gr.Slider(label="Perspectives (%)", minimum=0, maximum=100, value=50)
+                pct_us = gr.Slider(label="US (%)", minimum=0, maximum=100, value=30)
+                max_per_author = gr.Slider(label="Max articles per author", minimum=1, maximum=20, value=6, step=1)
+        results = gr.DataFrame(
+            label="Results",
+            type="array",
+            headers=["#", "Q", "%", "Date", "Authors", "Title"],
+            datatype=["number", "number", "number", "date", "str", "markdown"],
+            column_widths=["3.5%", "3.5%", "4%", "7%", "20%", "62%"],
+        )
+        gather.click(
+            fn=gather_articles,
+            inputs=[query, year, articles_per_year, pct_perspectives, pct_us, max_per_author],
+            outputs=[results],
+        )

article_embedding/app/lookup.py ADDED Viewed

	@@ -0,0 +1,11 @@

+import gradio as gr
+from article_embedding.couchdb import CouchDB
+def lookup_tab() -> None:
+    with gr.Tab("Lookup"):
+        doc_id = gr.Textbox(label="Document ID", placeholder="Enter a document ID")
+        button = gr.Button("Lookup")
+        doc_output = gr.JSON()
+        button.click(fn=CouchDB().get_doc, inputs=[doc_id], outputs=[doc_output])

article_embedding/app/model.py ADDED Viewed

	@@ -0,0 +1,71 @@

+import datetime
+import logging
+from typing import Any
+from pydantic import BaseModel
+log = logging.getLogger(__name__)
+class WorkflowData(BaseModel):
+    status: str | None = None
+    writers: list[str] | None = None
+    editors: list[str] | None = None
+    proofers: list[str] | None = None
+    reviewers: list[str] | None = None
+    proofingDeadline: str | None = None
+class Document(BaseModel):
+    _id: str
+    _rev: str
+    type: str | None = None
+    mimetype: str | None = None
+    title: str | None = None
+    language: str | None = None
+    workflowData: WorkflowData | None = None
+    path: str | None = None
+    name: str | None = None
+    created: int | None = None
+    creator: str | None = None
+    lastPublished: int | None = None
+    firstPublished: int | None = None
+    modified: datetime.datetime | None = None
+    modifier: str | None = None
+    published: datetime.datetime | None = None
+    authors: list[str] | None = None
+    content: str | None = None
+    contentAssets: list[str] | None = None
+    featuredImages: list[str] | None = None
+    keywords: list[str] | None = None
+    topics: list[str] | None = None
+    relatedAssets: list[str] | None = None
+    comments: bool | None = None
+    campaignConfigs: list[Any] | None = None
+    order: int | None = None
+    overline: str | None = None
+    translatedFrom: str | None = None
+    socialTitles: list[Any] | None = None
+    socialDescriptions: list[Any] | None = None
+    socialFeaturedImages: list[Any] | None = None
+    underline: str | None = None
+    template: str | None = None
+    description: str | None = None
+    suggestedImages: list[str] | None = None
+    publisher: str | None = None
+    @property
+    def url(self) -> str:
+        if not self.path:
+            log.warning(f"URL requested for pathless document: {self}.")
+        return public_url(self.path)
+    def has_topic(self, topic_id: str) -> bool:
+        return topic_id in self.topics if self.topics else False
+    def __str__(self) -> str:
+        return f"Document[id={self._id!r}, path={self.path!r}]"
+def public_url(path: str | None) -> str:
+    return f"https://www.wsws.org{path or ''}"

article_embedding/app/rag.py ADDED Viewed

	@@ -0,0 +1,134 @@

+import asyncio
+from typing import AsyncGenerator, Literal
+import gradio as gr
+from gradio import MessageDict
+from openai import AsyncOpenAI, BaseModel
+from openai.types.chat import ChatCompletionMessageParam
+from article_embedding.chunk import html_to_markdown
+from article_embedding.couchdb import CouchDB
+from article_embedding.query import Query
+from article_embedding.utils import env_str
+BackendId = Literal["OpenAI", "Ollama"]
+class BackendConfig(BaseModel):
+    base_url: str | None = None
+    api_key: str | None = None
+    @property
+    def async_client(self) -> AsyncOpenAI:
+        return AsyncOpenAI(api_key=self.api_key, base_url=self.base_url)
+BACKEND_CONFIG: dict[BackendId, BackendConfig] = {
+    "OpenAI": BackendConfig(),
+    "Ollama": BackendConfig(base_url=env_str("OLLAMA_URL"), api_key="Ollama"),
+}
+class RagModel(BaseModel):
+    backend_id: BackendId
+    model: str
+    @property
+    def backend(self) -> BackendConfig:
+        return BACKEND_CONFIG[self.backend_id]
+RAG_MODELS: dict[str, RagModel] = {
+    m.model: m
+    for m in [
+        RagModel(backend_id="OpenAI", model="gpt-4o"),
+        RagModel(backend_id="OpenAI", model="gpt-4o-mini"),
+        RagModel(backend_id="Ollama", model="qwen2.5:72b"),
+        RagModel(backend_id="Ollama", model="llama3.3:70b"),
+        RagModel(backend_id="Ollama", model="phi4"),
+    ]
+}
+QUERY_PROMPT = """
+Generate a short paragraph whose vector embedding will be used to find related articles using a vector database indexing the WSWS archive.
+Only output the paragraph without any additional text.
+{question}
+""".strip()  # noqa: E501
+CONTEXT_PROMPT = """
+Using the following sources, answer the user's question as precisely as possible.
+Only use the information provided in the sources.
+If not enough information is available, say you don't know.
+Include references to the sources in the answer by using just the article index and url, e.g. [3](URL).
+The answers should be exhaustive and cover as much as possible from the sources.
+{docs}
+""".strip()
+_couchdb = CouchDB()
+# From Gradio chat_interface.py:ChatInterface:_stream_fn
+async def chat_function(
+    message: str,  # | MultimodalPostprocess,
+    history_with_input: list[MessageDict],
+    model: str,
+) -> AsyncGenerator[str | list[MessageDict], None]:
+    _openai = RAG_MODELS[model].backend.async_client
+    _chat = _openai.chat.completions.create
+    messages: list[ChatCompletionMessageParam] = [{"role": h["role"], "content": h["content"]} for h in history_with_input]
+    async def generate_system_prompt(question: str) -> str:
+        msgs: list[ChatCompletionMessageParam] = messages + [{"role": "system", "content": QUERY_PROMPT.format(question=question)}]
+        response = await _chat(messages=msgs, model=model, stream=False)
+        query = response.choices[0].message.content or ""
+        result = Query(index="wsws-2").query(query, limit=40)
+        paths = []
+        for point in result.points:
+            doc = point.payload
+            if doc is None:
+                continue
+            path = doc["path"]
+            if path not in paths:
+                paths.append(path)
+            if len(paths) >= 8:
+                break
+        tasks = [_couchdb.get_doc(path) for path in paths]
+        docs: list[str] = []
+        for rix, task_result in enumerate(await asyncio.gather(*tasks, return_exceptions=True), start=1):
+            if isinstance(task_result, dict):
+                data = [
+                    f"Index: {rix}",
+                    f"URL: https://www.wsws.org{task_result["path"]}",
+                    task_result["title"] or "",
+                    task_result["description"] or "",
+                    html_to_markdown(task_result["content"]),
+                    "---",
+                ]
+                docs.append("\n".join(data))
+        return CONTEXT_PROMPT.format(docs="\n\n".join(docs))
+    system_prompt = await generate_system_prompt(message)
+    messages.append({"role": "system", "content": system_prompt})
+    messages.append({"role": "user", "content": message})
+    generator = await _chat(messages=messages, model=model, stream=True)
+    result: MessageDict = {"role": "assistant", "content": ""}
+    async for chunk in generator:
+        delta = chunk.choices[0].delta
+        if delta.content:
+            result["content"] += delta.content
+            yield result
+def rag_tab() -> None:
+    with gr.Tab("RAG"):
+        model = gr.Dropdown(value="gpt-4o", choices=list(RAG_MODELS), label="Model")
+        gr.ChatInterface(
+            fn=chat_function,
+            multimodal=False,
+            type="messages",
+            additional_inputs=[model],
+            chatbot=gr.Chatbot(height=800),
+        )

article_embedding/app/search.py ADDED Viewed

	@@ -0,0 +1,70 @@

+from datetime import datetime
+from typing import Any
+import gradio as gr
+from qdrant_client.models import Condition, Filter
+from article_embedding.query import Query, make_date_condition, make_topic_condition
+def search(
+    text: str,
+    date_from: datetime | None,
+    date_to: datetime | None,
+    topics: str,
+    limit: int,
+) -> tuple[list[list[str]], list[Any]]:
+    rv: list[list[str]] = []
+    topic_lines = [t for t in topics.strip().split("\n") if t]
+    should: list[Condition] = [make_topic_condition(t) for t in topic_lines if t[0].isalnum()]
+    must: list[Condition] = [make_topic_condition(t[1:]) for t in topic_lines if t.startswith("+")]
+    must_not: list[Condition] = [make_topic_condition(t[1:]) for t in topic_lines if t.startswith("!")]
+    if date_from or date_to:
+        must.append(make_date_condition(date_from=date_from, date_to=date_to))  # type: ignore
+    result = Query().query(
+        text,
+        query_filter=Filter(should=should, must=must, must_not=must_not),
+        limit=limit,
+    )
+    docs = []
+    for point in result.points:
+        doc = point.payload
+        assert doc is not None
+        docs.append(doc)
+        rv.append(
+            [
+                f"{point.score * 100:.1f}%",
+                datetime.fromtimestamp(doc["published"]).strftime("%Y/%m/%d"),
+                ", ".join(doc["authors"]),
+                f'[{doc["title"] or "N/A"}](https://www.wsws.org{doc["path"]})',
+            ]
+        )
+    return rv, docs
+def search_tab() -> None:
+    default_query = "The COVID winter wave, the emergence of the Delta variant and the January 6th coup"
+    with gr.Tab("Search"):
+        with gr.Row():
+            with gr.Column(scale=2):
+                query = gr.Textbox(
+                    label="Query",
+                    placeholder="Enter a query",
+                    value=default_query,
+                )
+            with gr.Column():
+                from_date = gr.DateTime(include_time=False, type="datetime", label="From", value="2021-01-01")
+                to_date = gr.DateTime(include_time=False, type="datetime", label="To", value="2021-05-01")
+                topics = gr.Textbox(label="Topics", placeholder="Enter topics")
+                limit = gr.Number(label="Limit", minimum=1, value=10)
+                button = gr.Button("Search")
+        results = gr.DataFrame(
+            label="Results",
+            type="array",
+            headers=["Score", "Published", "Authors", "Title"],
+            datatype=["number", "date", "str", "markdown"],
+        )
+        with gr.Accordion("Debug"):
+            json_output = gr.JSON()
+    button.click(fn=search, inputs=[query, from_date, to_date, topics, limit], outputs=[results, json_output])

article_embedding/app/ui.py ADDED Viewed

	@@ -0,0 +1,31 @@

+import asyncio
+import os
+import gradio as gr
+from dotenv import load_dotenv
+from article_embedding.app.gather import gather_tab
+from article_embedding.app.lookup import lookup_tab
+from article_embedding.app.rag import rag_tab
+from article_embedding.app.search import search_tab
+load_dotenv()
+async def amain(share: bool = False) -> None:
+    with gr.Blocks(fill_height=True) as app:  # noqa: SIM117
+        rag_tab()
+        gather_tab()
+        search_tab()
+        lookup_tab()
+    auth = (os.environ["GRADIO_USERNAME"], os.environ["GRADIO_PASSWORD"]) if share else None
+    app.queue().launch(inbrowser=True, share=share, auth=auth)
+def run_app(share: bool = False) -> None:
+    asyncio.run(amain(share))
+if __name__ == "__main__":
+    run_app()

article_embedding/benchmark.py ADDED Viewed

	@@ -0,0 +1,76 @@

+import json
+import logging
+import time
+from pathlib import Path
+from typing import Any
+import numpy as np
+import numpy.typing as npt
+import torch
+from article_embedding.embed import SentenceTransformerModel, StellaEmbedder
+log = logging.getLogger(__name__)
+async def fetch_documents() -> list[str]:
+    from article_embedding.couchdb import CouchDB
+    async for sorted_batch in CouchDB().changes(batch_size=128):
+        sorted_batch = sorted(sorted_batch, key=lambda x: x.get("_id"))
+        return [a["content"] for a in sorted_batch if a.get("type") == "article" and a.get("language") == "en"]
+    return []
+def process2(model: SentenceTransformerModel, documents: list[str], name: str) -> None:
+    model.embed(["Hello, world!"])  # Warmup
+    ts0 = time.time()
+    embeddings = model.embed(documents)
+    benchmark = time.time() - ts0
+    output_path = Path("data/embeddings.json")
+    save_embeddings(embeddings, output_path)
+    golden_path = Path(f"data/embeddings.{name}-golden.json")
+    if golden_path.exists():
+        golden_embeddings: Any = load_embeddings(golden_path)
+        similarities = model.model.similarity_pairwise(embeddings, golden_embeddings)
+        rms = torch.sqrt(torch.mean(similarities**2)).item()
+    else:
+        save_embeddings(embeddings, golden_path)
+        rms = 0.0
+    log.info(
+        "%s - RMS: %.2f. Latency: %.2f ms. Size: %d",
+        name,
+        rms,
+        benchmark / len(embeddings) * 1000,
+        len(embeddings[0]),
+    )
+def load_embeddings(path: Path) -> list[npt.NDArray[np.float64]]:
+    with path.open() as f:
+        return [np.array(json.loads(line)) for line in f.readlines()]
+def save_embeddings(embeddings: list[npt.NDArray[np.float64]], path: Path) -> None:
+    with path.open("w") as f:
+        for e in embeddings:
+            f.write(json.dumps(e.tolist()) + "\n")
+async def amain() -> None:
+    model = StellaEmbedder()
+    # model = NvEmbedder()
+    # model = JasperEmbedder()
+    # model.model.half()
+    documents = await fetch_documents()
+    process2(model, documents, "stella")
+    # documents = await fetch_documents()
+    # process2(model, documents, "stella")
+if __name__ == "__main__":
+    import asyncio
+    logging.basicConfig(level=logging.WARN)
+    log.setLevel(logging.DEBUG)
+    asyncio.run(amain())

article_embedding/chunk.py ADDED Viewed

	@@ -0,0 +1,25 @@

+from abc import ABC, abstractmethod
+from typing import cast
+import pandoc
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+class Chunker(ABC):
+    @abstractmethod
+    def chunk(self, text: str) -> list[str]: ...
+class LangchainChunker(Chunker):
+    def __init__(self, chunk_size: int = 576, chunk_overlap: int = 0) -> None:
+        self.splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
+    def chunk(self, text: str) -> list[str]:
+        md = html_to_markdown(text)
+        return self.splitter.split_text(md)
+def html_to_markdown(html: str) -> str:
+    doc = pandoc.read(html, format="html")
+    md = pandoc.write(doc, format="markdown-smart", options=["--wrap=none"])
+    return cast(str, md)

article_embedding/cli.py ADDED Viewed

	@@ -0,0 +1,24 @@

+import logging
+import click
+from click_help_colors import HelpColorsGroup
+from article_embedding.app import run_app
+@click.group(
+    cls=HelpColorsGroup,
+    help_headers_color="yellow",
+    help_options_color="green",
+)
+def cli() -> None:
+    pass
+@cli.command()
+@click.option("--share", is_flag=True, help="Share the app at Gradio.")
+def app(share: bool) -> None:
+    """Run the Gradio app."""
+    logging.basicConfig(level=logging.WARN, format="%(asctime)s %(levelname)s %(name)s - %(message)s")
+    logging.getLogger("article_embedding").setLevel(logging.INFO)
+    run_app(share)

article_embedding/constants.py ADDED Viewed

	@@ -0,0 +1,2 @@


1	+ PERSPECTIVE_TOPIC_ID = "105b9ec2-0cdd-480d-9c55-6f806fa7b9bc"
2	+ US_TOPIC_ID = "ea69254c-a976-4e67-a79c-843202a10e55"

article_embedding/couchdb.py ADDED Viewed

	@@ -0,0 +1,203 @@

+import contextlib
+import json
+import logging
+import os
+import re
+from abc import ABC, abstractmethod
+from typing import Any, AsyncGenerator
+import backoff
+import httpx
+import tqdm
+log = logging.getLogger(__name__)
+class Checkpoint(ABC):
+    @abstractmethod
+    def get(self) -> str | None: ...
+    @abstractmethod
+    def set(self, value: str) -> None: ...
+    @abstractmethod
+    def reset(self) -> None: ...
+class NullCheckpoint(Checkpoint):
+    def get(self) -> str | None:
+        return None
+    def set(self, value: str) -> None:
+        pass
+    def reset(self) -> None:
+        pass
+_NULL_CHECKPOINT = NullCheckpoint()
+class FileCheckpoint(Checkpoint):
+    def __init__(self, path: str) -> None:
+        self.path = path
+    def get(self) -> str | None:
+        try:
+            with open(self.path) as file:
+                return file.read().strip()
+        except FileNotFoundError:
+            return None
+    def set(self, value: str) -> None:
+        with open(self.path, "w") as file:
+            file.write(value)
+    def reset(self) -> None:
+        with contextlib.suppress(FileNotFoundError):
+            os.remove(self.path)
+class CouchDB:
+    def __init__(self) -> None:
+        self.client = self.make_client()
+        self.database = os.environ["COUCHDB_DB"]
+        self.path_view = f"/{self.database}/{os.environ["DOCS_PATH_VIEW"]}"
+    def __new__(cls) -> "CouchDB":
+        if not hasattr(cls, "_instance"):
+            cls._instance = super().__new__(cls)
+        return cls._instance
+    def make_client(self) -> httpx.AsyncClient:
+        url = os.environ["COUCHDB_URL"]
+        user = os.environ["COUCHDB_USER"]
+        password = os.environ["COUCHDB_PASSWORD"]
+        auth = {"name": user, "password": password}
+        async def on_backoff(details: Any) -> None:
+            response = await self.client.post("/_session", json=auth)
+            response.raise_for_status()
+        client = httpx.AsyncClient(base_url=url)
+        decorator = backoff.on_predicate(
+            backoff.expo,
+            predicate=lambda r: r.status_code == 401,
+            on_backoff=on_backoff,
+            max_tries=2,
+            factor=0,
+        )
+        client.get = decorator(client.get)  # type: ignore[method-assign]
+        return client
+    async def changes(self, *, batch_size: int, checkpoint: Checkpoint = _NULL_CHECKPOINT) -> AsyncGenerator[list[Any], None]:
+        since = checkpoint.get() or 0
+        params = {"since": since, "limit": batch_size, "include_docs": True}
+        while True:
+            response = await self.client.get(f"/{self.database}/_changes", params=params)
+            response.raise_for_status()
+            data = response.json()
+            yield [change["doc"] for change in data["results"]]
+            since = data["last_seq"]
+            assert isinstance(since, str)
+            params["since"] = since
+            checkpoint.set(since)
+            if data["pending"] == 0:
+                break
+    async def estimate_total_changes(self, *, checkpoint: Checkpoint = _NULL_CHECKPOINT) -> int:
+        since = checkpoint.get() or 0
+        params = {"since": since, "limit": 0}
+        response = await self.client.get(f"/{self.database}/_changes", params=params)
+        response.raise_for_status()
+        data = response.json()
+        return int(data["pending"]) + 1
+    async def get_doc_by_id(self, doc_id: str) -> Any:
+        try:
+            response = await self.client.get(f"/{self.database}/{doc_id}")
+            if response.status_code == 404:
+                return None
+            response.raise_for_status()
+            return response.json()
+        except Exception as e:
+            log.error("Error fetching document by ID", exc_info=e)
+            return None
+    async def get_doc_by_path(self, path: str) -> Any:
+        try:
+            params = {
+                "limit": "1",
+                "key": json.dumps(path),
+                "include_docs": "true",
+            }
+            response = await self.client.get(self.path_view, params=params)
+            response.raise_for_status()
+            data = response.json()
+            rows = data["rows"]
+            if not rows:
+                return None
+            return rows[0]["doc"]
+        except Exception as e:
+            logging.error("Error fetching document by path", exc_info=e)
+            return None
+    async def get_doc(self, id_or_path: str) -> Any:
+        uuids = extract_doc_ids(id_or_path)
+        for uuid in uuids:
+            doc = await self.get_doc_by_id(uuid)
+            if doc:
+                return doc
+        path = extract_doc_path(id_or_path)
+        if path:
+            return await self.get_doc_by_path(path)
+        return None
+UUID_PATTERN = re.compile(r"[0-9a-f]{8}-[0-9a-f]{4}-[1-5][0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}")
+def extract_doc_ids(s: str) -> list[str]:
+    return UUID_PATTERN.findall(s)
+def extract_doc_path(s: str) -> str | None:
+    if not s.endswith(".html"):
+        return None
+    if s.startswith("/"):
+        return s
+    if "://" in s:
+        s = s.split("://", 1)[1]
+        if "/" in s:
+            return "/" + s.split("/", 1)[1]
+    return None
+if __name__ == "__main__":
+    async def main() -> None:
+        db = CouchDB()
+        checkpoint = FileCheckpoint(".checkpoint")
+        total = await db.estimate_total_changes(checkpoint=checkpoint)
+        with tqdm.tqdm(total=total) as pbar:
+            async for docs in db.changes(batch_size=40, checkpoint=checkpoint):
+                for doc in docs:
+                    kind = doc.get("type")
+                    if kind == "article":
+                        _id = doc["_id"]
+                        language = doc["language"]
+                        path = doc["path"]
+                        path = os.path.basename(path)
+                        pbar.desc = f"{_id}: {kind} {language} {path}"
+                    else:
+                        pbar.desc = f"{kind}"
+                    pbar.update(1)
+    import asyncio
+    from dotenv import load_dotenv
+    load_dotenv()
+    asyncio.run(main())

article_embedding/embed.py ADDED Viewed

	@@ -0,0 +1,86 @@

+from abc import ABC, abstractmethod
+from typing import Any
+from fastembed import TextEmbedding  # type: ignore
+from qdrant_client import QdrantClient
+from qdrant_client.models import Distance, VectorParams
+from sentence_transformers import SentenceTransformer
+class EmbeddingModel(ABC):
+    @abstractmethod
+    def embed(self, documents: list[str]) -> Any: ...
+    @abstractmethod
+    def get_vectors_config(self) -> VectorParams: ...
+class SentenceTransformerModel(EmbeddingModel):
+    model: SentenceTransformer
+    def embed(self, documents: list[str]) -> Any:
+        return self.model.encode(documents, normalize_embeddings=True)
+    def get_vectors_config(self) -> VectorParams:
+        dimensions = self.model.get_sentence_embedding_dimension()
+        assert dimensions is not None
+        return VectorParams(
+            size=dimensions,
+            distance=Distance.COSINE,
+        )
+class StellaEmbedder(SentenceTransformerModel):
+    def __init__(self) -> None:
+        self.model = SentenceTransformer(
+            "dunzhang/stella_en_400M_v5",
+            trust_remote_code=True,
+            config_kwargs={"use_memory_efficient_attention": False, "unpad_inputs": False},
+        ).to("mps")
+class NvEmbedder(SentenceTransformerModel):
+    def __init__(self) -> None:
+        self.model = SentenceTransformer(
+            "nvidia/NV-Embed-V2",
+            trust_remote_code=True,
+        ).to("mps")
+        self.model.max_seq_length = 32768
+        self.model.tokenizer.padding_side = "right"
+    def embed(self, documents: list[str]) -> Any:
+        processed = [d + self.model.tokenizer.eos_token for d in documents]
+        return self.model.encode(processed, normalize_embeddings=True)
+class JasperEmbedder(SentenceTransformerModel):
+    def __init__(self, use_sdpa: bool = True) -> None:
+        model_kwargs = {"attn_implementation": "sdpa"} if use_sdpa else {}
+        self.model = SentenceTransformer(
+            "infgrad/jasper_en_vision_language_v1",
+            trust_remote_code=True,
+            model_kwargs=model_kwargs,
+            config_kwargs={
+                "is_text_encoder": True,
+                "vector_dim": 1024,
+            },
+        ).to("mps")
+        self.model.max_seq_length = 1024
+    def embed(self, documents: list[str]) -> Any:
+        processed = [d + self.model.tokenizer.eos_token for d in documents]
+        return self.model.encode(processed, normalize_embeddings=True)
+class FastEmbedModel(EmbeddingModel):
+    def __init__(self, model_name: str = "BAAI/bge-small-en-v1.5") -> None:
+        self.model = TextEmbedding(model_name=model_name)
+    def embed(self, documents: list[str]) -> Any:
+        return self.model.embed(documents)
+    def get_vectors_config(self) -> VectorParams:
+        size, distance = QdrantClient._get_model_params(model_name=self.model.model_name)
+        return VectorParams(
+            size=size,
+            distance=distance,
+        )

article_embedding/loader.py ADDED Viewed

	@@ -0,0 +1,146 @@

+import contextlib
+import os
+import uuid
+from dataclasses import dataclass
+from typing import Any, Generator
+import tqdm
+from qdrant_client import QdrantClient
+from qdrant_client.http.exceptions import UnexpectedResponse
+from qdrant_client.models import Batch
+from article_embedding.chunk import Chunker, LangchainChunker
+from article_embedding.couchdb import Checkpoint, CouchDB, FileCheckpoint
+from article_embedding.embed import EmbeddingModel, StellaEmbedder
+class Loader:
+    def __init__(
+        self,
+        *,
+        collection_name: str,
+        embedding_model: EmbeddingModel,
+        checkpoint: Checkpoint,
+        chunker: Chunker | None = None,
+    ) -> None:
+        self.collection_name = collection_name
+        self.embedding_model = embedding_model
+        self.checkpoint = checkpoint
+        self.chunker = chunker
+        self.qdrant = QdrantClient(os.getenv("QDRANT_URL"))
+        self.couchdb = CouchDB()
+    @property
+    def starting_from_scratch(self) -> bool:
+        return self.checkpoint.get() is None
+    def recreate_collection(self) -> None:
+        with contextlib.suppress(UnexpectedResponse):
+            self.qdrant.delete_collection(self.collection_name)
+        self.qdrant.create_collection(
+            collection_name=self.collection_name,
+            vectors_config=self.embedding_model.get_vectors_config(),
+        )
+    async def load(self, batch_size: int = 64) -> None:
+        if self.starting_from_scratch:
+            self.recreate_collection()
+        total = await self.couchdb.estimate_total_changes(checkpoint=self.checkpoint)
+        with tqdm.tqdm(total=total) as pbar:
+            async for docs in self.couchdb.changes(batch_size=batch_size, checkpoint=self.checkpoint):
+                points = self.embed(docs)
+                if points:
+                    self.qdrant.upsert(
+                        collection_name=self.collection_name,
+                        points=points,
+                        wait=False,
+                    )
+                pbar.update(len(docs))
+    def embed(self, docs: list[Any]) -> Batch | None:
+        def article_filter(doc: Any) -> bool:
+            return doc.get("type") == "article" and doc.get("language") == "en" and doc.get("content") is not None
+        def to_payload(doc: Any) -> Any:
+            doc.pop("content")
+            return doc
+        def chunk(doc: dict[str, Any]) -> Generator[dict[str, Any], None, None]:
+            assert self.chunker is not None
+            content = doc["content"]
+            for ix, chunk in enumerate(self.chunker.chunk(content)):
+                copy = doc.copy()
+                copy["docId"] = doc["_id"]
+                copy["_id"] = str(uuid.uuid3(_NS_DOCUMENT, chunk))
+                copy["content"] = chunk
+                copy["chunkIndex"] = ix
+                copy["chunk"] = chunk
+                yield copy
+        filtered = [d for d in docs if article_filter(d)]
+        chunked = [c for d in filtered for c in chunk(d)] if self.chunker else filtered
+        if not chunked:
+            return None
+        ids = [d["_id"] for d in chunked]
+        docs = [d["content"] for d in chunked]
+        payloads = [to_payload(d) for d in chunked]
+        return Batch(
+            ids=ids,
+            vectors=self.embedding_model.embed(docs),
+            payloads=payloads,
+        )
+_NS_DOCUMENT = uuid.UUID("a07b4024-c09b-11ef-9e49-36d763a413b8")
+def uuid_for_content(*content: Any) -> str:
+    import hashlib
+    import uuid
+    hash = hashlib.md5("".join(map(str, content)).encode()).hexdigest()
+    return str(uuid.UUID(hash[:32]))
+if __name__ == "__main__":
+    stella = StellaEmbedder()
+    @dataclass
+    class Store:
+        collection_name: str
+        embedding_model: EmbeddingModel = stella
+        chunker: Chunker | None = None
+        batch_size: int = 64
+        @property
+        def checkpoint(self) -> str:
+            return f".{self.collection_name}.checkpoint"
+    class StoreManager:
+        def __init__(self) -> None:
+            self.stores: dict[str, Store] = {}
+        def add(self, store: Store) -> None:
+            self.stores[store.collection_name] = store
+        async def load(self, name: str) -> None:
+            store = self.stores[name]
+            checkpoint = FileCheckpoint(store.checkpoint)
+            loader = Loader(
+                collection_name=store.collection_name,
+                embedding_model=store.embedding_model,
+                checkpoint=checkpoint,
+                chunker=store.chunker,
+            )
+            await loader.load(batch_size=store.batch_size)
+        async def load_all(self) -> None:
+            for name in self.stores:
+                await self.load(name)
+    import asyncio
+    store_manager = StoreManager()
+    store_manager.add(Store(collection_name="wsws"))
+    store_manager.add(Store(collection_name="wsws-2", chunker=LangchainChunker(), batch_size=8))
+    asyncio.run(store_manager.load_all())

article_embedding/modal_app.py ADDED Viewed

	@@ -0,0 +1,153 @@

+import json
+import logging
+import time
+from pathlib import Path
+from typing import Any, Callable
+import modal
+import numpy as np
+import numpy.typing as npt
+import torch
+from sentence_transformers import SentenceTransformer
+from article_embedding.embed import SentenceTransformerModel, StellaEmbedder
+log = logging.getLogger(__name__)
+def load_model() -> SentenceTransformer:
+    return SentenceTransformer(
+        "dunzhang/stella_en_400M_v5",
+        trust_remote_code=True,
+        config_kwargs={"use_memory_efficient_attention": False, "unpad_inputs": False},
+    )
+image = modal.Image.debian_slim(python_version="3.12").pip_install(["sentence-transformers"]).run_function(load_model)
+app = modal.App("embedding", image=image)
+@app.cls(gpu="A10G")
+class ModalEmbedder:
+    @modal.enter()
+    def setup(self) -> None:
+        logging.basicConfig(level=logging.WARN)
+        log.setLevel(logging.DEBUG)
+        if torch.cuda.is_available():
+            device = "cuda"
+        elif torch.backends.mps.is_available():
+            device = "mps"
+        else:
+            device = "cpu"
+        self.model = load_model().to(device)
+        log.info("Model loaded on %s", device)
+    @modal.method()
+    def embed(self, documents: list[str]) -> Any:
+        return self.model.encode(documents)
+async def fetch_documents() -> list[str]:
+    from article_embedding.couchdb import CouchDB
+    async for sorted_batch in CouchDB().changes(batch_size=256):
+        sorted_batch = sorted(sorted_batch, key=lambda x: x.get("_id"))
+        return [a["content"] for a in sorted_batch if a.get("type") == "article" and a.get("language") == "en"]
+    return []
+def process(func: Callable[[list[str]], Any], documents: list[str], name: str) -> None:
+    func(["Hello, world!"])  # Warmup
+    ts0 = time.time()
+    embeddings = func(documents)
+    benchmark = time.time() - ts0
+    output_path = Path("data/embeddings.json")
+    golden_path = Path(f"data/embeddings.{name}-golden.json")
+    save_embeddings(embeddings, output_path)
+    cosine_distance, rms = compare_embeddings(embeddings, golden_path)
+    log.info(
+        "%s - MCS: %.2f. RMS: %.2f. Latency: %.2f ms. Size: %d",
+        name,
+        cosine_distance,
+        rms,
+        benchmark / len(embeddings) * 1000,
+        len(embeddings[0]),
+    )
+def process2(model: SentenceTransformerModel, documents: list[str], name: str) -> None:
+    model.embed(["Hello, world!"])  # Warmup
+    ts0 = time.time()
+    embeddings = model.embed(documents)
+    benchmark = time.time() - ts0
+    output_path = Path("data/embeddings.json")
+    save_embeddings(embeddings, output_path)
+    golden_path = Path(f"data/embeddings.{name}-golden.json")
+    if golden_path.exists():
+        golden_embeddings: Any = load_embeddings(golden_path)
+        similarities = model.model.similarity_pairwise(embeddings, golden_embeddings)
+        rms = torch.sqrt(torch.mean(similarities**2))
+    else:
+        save_embeddings(embeddings, golden_path)
+        rms = torch.zeros([0])
+    log.info(
+        "%s - RMS: %.2f. Latency: %.2f ms. Size: %d",
+        name,
+        rms,
+        benchmark / len(embeddings) * 1000,
+        len(embeddings[0]),
+    )
+def load_embeddings(path: Path) -> list[npt.NDArray[np.float64]]:
+    with path.open() as f:
+        return [np.array(json.loads(line)) for line in f.readlines()]
+def save_embeddings(embeddings: list[npt.NDArray[np.float64]], path: Path) -> None:
+    with path.open("w") as f:
+        for e in embeddings:
+            f.write(json.dumps(e.tolist()) + "\n")
+def compare_embeddings(embeddings: list[npt.NDArray[np.float64]], golden_path: Path) -> tuple[float, float]:
+    if not golden_path.exists():
+        save_embeddings(embeddings, golden_path)
+        return 0.0, 0.0
+    with golden_path.open() as f:
+        golden_embeddings = [np.array(json.loads(line)) for line in f.readlines()]
+    np_embeddings = np.array(embeddings)
+    np_golden_embeddings = np.array(golden_embeddings)
+    rms = np.sqrt(np.mean((np_embeddings - np_golden_embeddings) ** 2))
+    dot_products = np.einsum("ij,ij->i", np_embeddings, np_golden_embeddings)
+    norms = np.linalg.norm(np_embeddings, axis=1) * np.linalg.norm(np_golden_embeddings, axis=1)
+    cosine_similarities = dot_products / norms
+    return np.mean(cosine_similarities), np.mean(rms)
+@app.local_entrypoint()
+async def modal_amain() -> None:
+    logging.basicConfig(level=logging.WARN)
+    log.setLevel(logging.DEBUG)
+    embedder = ModalEmbedder()
+    documents = await fetch_documents()
+    process(embedder.embed.remote, documents, "modal")
+async def amain() -> None:
+    model = StellaEmbedder()
+    # model = NvEmbedder()
+    # model = JasperEmbedder()
+    # model.model.half()
+    documents = await fetch_documents()
+    process2(model, documents, "stella")
+if __name__ == "__main__":
+    import asyncio
+    logging.basicConfig(level=logging.WARN)
+    log.setLevel(logging.DEBUG)
+    asyncio.run(amain())

article_embedding/query.py ADDED Viewed

	@@ -0,0 +1,117 @@

+import os
+import warnings
+from datetime import datetime
+from typing import Any, Optional
+from qdrant_client import QdrantClient
+from qdrant_client.http.models import QueryResponse
+from qdrant_client.models import FieldCondition, Filter, MatchValue, Range
+from article_embedding.embed import StellaEmbedder
+warnings.simplefilter(action="ignore", category=FutureWarning)
+def as_timestamp(date: datetime | str) -> float:
+    if isinstance(date, datetime):
+        return date.timestamp()
+    return datetime.strptime(date, "%Y-%m-%d").timestamp()
+def make_date_condition(
+    *, field: str = "published", date_from: datetime | str | None = None, date_to: datetime | str | None = None
+) -> FieldCondition | None:
+    kwargs = {}
+    if date_from:
+        kwargs["gte"] = as_timestamp(date_from)
+    if date_to:
+        kwargs["lt"] = as_timestamp(date_to)
+    if kwargs:
+        return FieldCondition(key=field, range=Range(**kwargs))
+    return None
+def make_topic_condition(topic_id: str) -> FieldCondition:
+    return FieldCondition(key="topics[]", match=MatchValue(value=topic_id))
+class Query:
+    _instance: Optional["Query"] = None
+    _embedding_model_instance: Optional[StellaEmbedder] = None
+    def __init__(self, index: str = "wsws") -> None:
+        self.embedding_model = Query.embedding_model_singleton()
+        self.qdrant = QdrantClient(os.environ["QDRANT_URL"])
+        self.index = index
+    @staticmethod
+    def embedding_model_singleton() -> StellaEmbedder:
+        if Query._embedding_model_instance is None:
+            Query._embedding_model_instance = StellaEmbedder()
+        return Query._embedding_model_instance
+    @staticmethod
+    def singleton() -> "Query":
+        if Query._instance is None:
+            Query._instance = Query()
+        return Query._instance
+    def embed(self, query: str) -> Any:
+        return self.embedding_model.embed([query])[0]
+    def query(
+        self,
+        query: str,
+        query_filter: Filter | None = None,
+        limit: int = 10,
+    ) -> QueryResponse:
+        vector = self.embedding_model.embed([query])[0]
+        return self.qdrant.query_points(self.index, query=vector, query_filter=query_filter, limit=limit)
+if __name__ == "__main__":
+    import gspread
+    from dotenv import load_dotenv
+    from gspread.utils import ValueInputOption
+    data = [
+        ("2021-01-01", "2021-05-01", "The COVID winter wave, the emergence of the Delta variant and the January 6th coup"),
+        (
+            "2021-05-01",
+            "2021-09-01",
+            "The COVID vaccine rollout, Biden declaring independence from COVID while the Delta wave continues",
+        ),
+        (
+            "2021-09-01",
+            "2022-01-01",
+            "The emergence of the COVID Omicron variant and the embrace of herd immunity by the ruling class",
+        ),
+    ]
+    load_dotenv()
+    query = Query()
+    rows: list[list[str]] = []
+    for date_from, date_to, sentence in data:
+        result = query.query(
+            sentence,
+            query_filter=Filter(should=make_date_condition(date_from=date_from, date_to=date_to)),
+        )
+        rows.append([sentence])
+        for point in result.points:
+            doc = point.payload
+            assert doc is not None
+            print(f'{point.score * 100:.1f}% https://www.wsws.org{doc["path"]} - {doc["title"]}')
+            rows.append(
+                [
+                    f"{point.score * 100:.1f}%",
+                    datetime.fromtimestamp(doc["published"]).strftime("%Y/%m/%d"),
+                    ", ".join(doc["authors"]),
+                    f'=hyperlink("https://www.wsws.org{doc["path"]}", "{doc["title"]}")',
+                ]
+            )
+        rows.append([])
+    gc = gspread.auth.oauth(credentials_filename=os.environ["GOOGLE_CREDENTIALS"])
+    sh = gc.open("COVID-19 Compilation")
+    ws = sh.get_worksheet(0)
+    ws.append_rows(rows, value_input_option=ValueInputOption.user_entered)

article_embedding/replicate.py ADDED Viewed

	@@ -0,0 +1,21 @@

+from qdrant_client import QdrantClient
+from article_embedding.utils import env_str
+remote_url = env_str("CLUSTER_1_ENDPOINT")
+remote_key = env_str("CLUSTER_1_API_KEY")
+collection_name = "wsws-2"
+local_url = env_str("QDRANT_URL")
+# model = StellaEmbedder()
+cloud_client = QdrantClient(url=remote_url, api_key=remote_key)
+# cloud_client.recreate_collection(
+#     collection_name=collection_name,
+#     vectors_config=model.get_vectors_config(),
+# )
+local_client = QdrantClient(local_url)
+local_client.migrate(cloud_client, collection_names=[collection_name])

article_embedding/retrieval.py ADDED Viewed

	@@ -0,0 +1,330 @@

+import logging
+import math
+from typing import Generator, Literal
+from pydantic import BaseModel
+from qdrant_client.models import Filter, ScoredPoint
+from article_embedding.app.model import Document
+from article_embedding.constants import PERSPECTIVE_TOPIC_ID, US_TOPIC_ID
+from article_embedding.query import Query, make_date_condition
+log = logging.getLogger(__name__)
+class Candidate(BaseModel):
+    doc_id: str
+    score: float
+    query_id: int
+    meta: Document
+    @classmethod
+    def from_point(cls, query_id: int, point: ScoredPoint) -> "Candidate":
+        assert isinstance(point.id, str)
+        return cls(doc_id=point.id, score=point.score, query_id=query_id, meta=Document.model_validate(point.payload))
+CandidatePool = dict[str, Candidate]
+Candidates = list[Candidate]
+Fractions = tuple[float, float]
+AuthorCounts = dict[str, int]
+# Qdrant retrieval function (assume exists)
+def qdrant_search(query_id: int, query: str, start_date: str, end_date: str, top_k: int) -> list[Candidate]:
+    # returns a list of candidates (doc_id, relevance_score, metadata)
+    # metadata includes: date, author, topics, etc.
+    query_filter = Filter(must=make_date_condition(date_from=start_date, date_to=end_date))
+    result = Query.singleton().query(query, query_filter=query_filter, limit=top_k)
+    return [Candidate.from_point(query_id, point) for point in result.points]
+def compute_count(docs: Candidates, *topic_ids: str) -> int:
+    return sum(1 if any(d.meta.has_topic(t) for t in topic_ids) else 0 for d in docs)
+def compute_fractions(docs: Candidates) -> Fractions:
+    oped_count = sum(d.meta.has_topic(PERSPECTIVE_TOPIC_ID) for d in docs)
+    us_count = sum(d.meta.has_topic(US_TOPIC_ID) for d in docs)
+    total = len(docs)
+    return (oped_count / total, us_count / total)
+def compute_author_counts(docs: Candidates) -> AuthorCounts:
+    counts: AuthorCounts = {}
+    for d in docs:
+        for author in d.meta.authors or []:
+            counts[author] = counts.get(author, 0) + 1
+    return counts
+Category = tuple[bool, bool]
+SwapIndices = tuple[int, int]
+SwapCandidates = tuple[Candidate, Candidate]
+Sign = Literal[-1, 0, 1]
+Evaluation = Literal["rejected", "accepted", "desired"]
+def sign(n: int) -> Sign:
+    return -1 if n < 0 else 0 if n == 0 else 1
+class Categories(BaseModel):
+    accepted: list[Sign] = []
+    desired: list[Sign] = []
+    def evaluate_delta(self, delta: int) -> Evaluation:
+        if sign(delta) in self.accepted:
+            return "accepted"
+        if sign(delta) in self.desired:
+            return "desired"
+        return "rejected"
+def seek_constraints(
+    selected: Candidates,
+    available: Candidates,
+    oped_target: int,
+    us_target: int,
+    max_per_author: int,
+    tolerance: int,
+) -> bool:
+    # Swaps candidates from `selected` with candidates from `available` with the goal of
+    # meeting the target counts for OpEd and US topics, while respecting the max per author constraint.
+    # and favoring a more balanced author count distribution (i.e. closer to the mean).
+    # It is a precondition that the author counts in `selected` do not already exceed the limit.
+    # The function will attempt to adjust the selection to meet the target counts approximately.
+    # The return value indicates if the algorithm was able to meet the constraints.
+    # This is a heuristic approach, i.e. it may not always find the optimal solution.
+    # It swaps documents from the tail of `selected` with appropriate replacements from
+    # the head of `available`.
+    # The max per author constraint is a hard constraint, i.e. it may not be exceeded.
+    # The algorithm also attempts to move the author count means closer to the mean.
+    # The count constraints are soft constraints, i.e. they allow some tolerance.
+    # The algorithm finishes when the counts meet the target within the tolerance.
+    # A band of ±2 is considered a hard constraint, ±1 is a soft constraint, and 0 is on target.
+    def compute_band(current: int, target: int, tolerance: int) -> int:
+        if current < target - tolerance:
+            return -2
+        if current == target - tolerance:
+            return -1
+        if current < target + tolerance:
+            return 0
+        if current == target + tolerance:
+            return +1
+        return +2
+    # Determine current counts and bands
+    oped_current = compute_count(selected, PERSPECTIVE_TOPIC_ID)
+    us_current = compute_count(selected, US_TOPIC_ID)
+    author_counts = compute_author_counts(selected)
+    author_mean_numerator = sum(author_counts.values())
+    author_mean_denominator = len(author_counts) or 1
+    # It is a precondition that the author counts in `selected` do not exceed the limit
+    if sum(1 for count in author_counts.values() if count > max_per_author):
+        raise ValueError("Precondition violated: author counts exceed the limit")
+    def iterate_swap_indices() -> Generator[SwapIndices, None, None]:
+        for dropped_index in range(len(selected) - 1, -1, -1):
+            for picked_index in range(len(available)):
+                yield (picked_index, dropped_index)
+    def compute_topic_delta(swap_candidates: SwapCandidates, topic_id: str) -> int:
+        return sum(d if c.meta.has_topic(topic_id) else 0 for d, c in zip([1, -1], swap_candidates))
+    def evaluate_author_change(author: str, delta: int, mean: float) -> int:
+        count = author_counts.get(author, 0)
+        old_distance = abs(count - mean)
+        new_distance = abs(count + delta - mean)
+        return int(math.copysign(1, old_distance - new_distance))
+    def evaluate_author_exchange(swap_candidates: SwapCandidates, mean: float) -> Evaluation:
+        points = 0
+        # An exchange gets points for moving author counts closer to the mean
+        picked, dropped = swap_candidates
+        for author_picked in picked.meta.authors or []:
+            if author_counts.get(author_picked, 0) >= max_per_author:
+                return "rejected"
+            points += evaluate_author_change(author_picked, +1, mean)
+        for author_dropped in dropped.meta.authors or []:
+            points += evaluate_author_change(author_dropped, -1, mean)
+        return "desired" if points > 0 else "accepted"
+    def update_author_counts(swap_candidates: SwapCandidates) -> int:
+        author_mean_numerator_delta = 0
+        picked, dropped = swap_candidates
+        for author in picked.meta.authors or []:
+            author_mean_numerator_delta += 1
+            author_counts[author] = author_counts.get(author, 0) + 1
+        for author in dropped.meta.authors or []:
+            author_mean_numerator_delta -= 1
+            new_count = author_counts.pop(author) - 1
+            if new_count:
+                author_counts[author] = new_count
+        return author_mean_numerator_delta
+    # Bands
+    # -2 - pick +, drop -
+    # -1 - pick +, drop -, pick +, drop +, pick -, drop -
+    #  0 - all
+    # +1 - pick -, drop +, pick +, drop +, pick -, drop -
+    # +2 - pick -, drop +
+    CATEGORIES = [
+        Categories(desired=[1]),
+        Categories(desired=[1], accepted=[0]),
+        Categories(accepted=[1, 0, -1]),
+        Categories(desired=[-1], accepted=[0]),
+        Categories(desired=[-1]),
+    ]
+    # We do multiple passes, in each successive pass we try to improve fewer constraints
+    swap_attempt = 0
+    swap_count = 0
+    for desired_target in range(3, 0, -1):
+        current_pass = 4 - desired_target
+        # Iterate all possible swaps
+        for picked_index, dropped_index in iterate_swap_indices():
+            swap_attempt += 1
+            swap_candidates = (available[picked_index], selected[dropped_index])
+            oped_band = compute_band(oped_current, oped_target, tolerance)
+            us_band = compute_band(us_current, us_target, tolerance)
+            author_mean_denominator = len(author_counts) or 1
+            author_mean = author_mean_numerator / author_mean_denominator
+            if log.isEnabledFor(logging.INFO):
+                author_variance = sum((count - author_mean) ** 2 for count in author_counts.values()) / author_mean_denominator
+            if log.isEnabledFor(logging.DEBUG):
+                # fmt: off
+                log.debug("%d/%d - %d Eval %d <-> %d: oped=%d (%d %d), us=%d (%d %d), author µ=%.2f σ=%.2f",
+                    swap_count, swap_attempt, current_pass, picked_index, dropped_index,
+                    oped_current, oped_target, oped_band, us_current, us_target, us_band, author_mean, author_variance,
+                ) # fmt: on
+            # Check if constraints are met
+            if abs(oped_band) <= 1 and abs(us_band) <= 1:
+                # fmt: off
+                log.info("%d/%d - %d Constraints met: oped=%d (%d %d), us=%d (%d %d), author µ=%.2f σ=%.2f",
+                    swap_count, swap_attempt, current_pass,
+                    oped_current, oped_target, oped_band, us_target, us_target, us_band, author_mean, author_variance,
+                ) # fmt: on
+                return True
+            oped_delta = compute_topic_delta(swap_candidates, PERSPECTIVE_TOPIC_ID)
+            category = CATEGORIES[oped_band + 2]
+            oped_evaluation = category.evaluate_delta(oped_delta)
+            if oped_evaluation == "rejected":
+                log.debug("%d/%d - %d Swap rejected by OpEd constraints", swap_count, swap_attempt, current_pass)
+                continue
+            us_delta = compute_topic_delta(swap_candidates, US_TOPIC_ID)
+            category = CATEGORIES[us_band + 2]
+            us_evaluation = category.evaluate_delta(us_delta)
+            if us_evaluation == "rejected":
+                log.debug("%d/%d - %d Swap rejected by US constraints", swap_count, swap_attempt, current_pass)
+                continue
+            author_evaluation = evaluate_author_exchange(swap_candidates, author_mean)
+            if author_evaluation == "rejected":
+                log.debug("%d/%d - %d Swap rejected by author constraints", swap_count, swap_attempt, current_pass)
+                continue
+            desired_count = sum(1 for f in [oped_evaluation, us_evaluation, author_evaluation] if f == "desired")
+            if desired_count < desired_target:
+                log.debug("%d/%d - %d Swap rejected by desired count", swap_count, swap_attempt, current_pass)
+                continue
+            swap_count += 1
+            author_mean_numerator_delta = update_author_counts(swap_candidates)
+            oped_current += oped_delta
+            us_current += us_delta
+            author_mean_numerator += author_mean_numerator_delta
+            picked, dropped = swap_candidates
+            selected.pop(dropped_index)
+            available.pop(picked_index)
+            selected.append(picked)
+            available.append(dropped)
+            if log.isEnabledFor(logging.INFO):
+                author_variance = sum((count - author_mean) ** 2 for count in author_counts.values()) / author_mean_denominator
+                # fmt: off
+                log.info("%d/%d - %d Swapped %d <-> %d: oped=%d (%d %d), us=%d (%d %d), author µ=%.2f σ=%.2f",
+                    swap_count, swap_attempt, current_pass, picked_index, dropped_index,
+                    oped_current, oped_target, oped_band, us_current, us_target, us_band, author_mean, author_variance,
+                ) # fmt: on
+    log.warning("No more improvements possible")
+    return False
+def retrieve_documents(
+    *,
+    queries: list[str],
+    start_date: str,
+    end_date: str,
+    number_of_articles: int,
+    op_ed_ratio: float,
+    us_ratio: float,
+    max_per_author: int,
+) -> Candidates:
+    K = 5 * number_of_articles  # retrieve 5x the needed documents from each query
+    # 1. Retrieve candidates for each query
+    candidate_pool: CandidatePool = {}
+    for query_id, q in enumerate(queries):
+        results = qdrant_search(query_id, q, start_date, end_date, K)
+        for candidate in results:
+            if candidate.doc_id not in candidate_pool:
+                candidate_pool[candidate.doc_id] = candidate
+            else:
+                # If the doc appears multiple times, keep the higher score
+                if candidate.score > candidate_pool[candidate.doc_id].score:
+                    candidate_pool[candidate.doc_id] = candidate
+    # 2. Sort candidates by relevance
+    sorted_candidates = sorted(candidate_pool.values(), key=lambda c: c.score, reverse=True)
+    # 3. Initial selection: take top M documents by relevance respecting author limits
+    available: Candidates = []
+    skipped: Candidates = []
+    author_counts: dict[str, int] = {}
+    for c in sorted_candidates:
+        skip = False
+        for author in c.meta.authors or []:
+            author_counts[author] = author_counts.get(author, 0) + 1
+            if author_counts[author] > max_per_author:
+                skip = True
+                break
+        if skip:
+            skipped.append(c)
+        else:
+            available.append(c)
+    if len(available) < number_of_articles:
+        raise ValueError("Not enough documents to meet the max_per_author constraint")
+    available += skipped
+    # 4. Attempt adjustments to meet constraints approximately
+    selected = available[:number_of_articles]
+    available = available[number_of_articles:]
+    seek_constraints(
+        selected=selected,
+        available=available,
+        oped_target=int(op_ed_ratio * number_of_articles),
+        us_target=int(us_ratio * number_of_articles),
+        max_per_author=max_per_author,
+        tolerance=int(0.05 * number_of_articles),
+    )
+    # selected now contains the chosen documents meeting date range, maximizing relevance,
+    # and approximately meeting the OpEd/Non-OpEd, US/International fractions, as well as
+    # respecting the max_per_author constraint.
+    return selected

article_embedding/sheets.py ADDED Viewed

	@@ -0,0 +1,7 @@

+import os
+import gspread
+gc = gspread.auth.oauth(credentials_filename=os.environ["GOOGLE_CREDENTIALS"])
+sh = gc.open("COVID-19 Compilation")
+print(sh.sheet1.get("A1"))

article_embedding/utils.py ADDED Viewed

	@@ -0,0 +1,76 @@

+"""Assorted utils."""
+import json
+import os
+from types import EllipsisType
+from typing import Any
+import yaml
+_FALSEY_STRINGS = ["", "0", "false", "no", "off", "n"]
+def is_truthy(value: Any) -> bool:
+    """Check if a value is truthy."""
+    if isinstance(value, str):
+        return value.lower() not in _FALSEY_STRINGS
+    return bool(value)
+def env_data(env_name: str, default: Any = ...) -> Any:
+    """Load a YAML/JSON from an environment variable."""
+    value = os.environ.get(env_name)
+    if value is None:
+        if default is ...:
+            raise ValueError(f"Missing environment variable {env_name}")
+        return default
+    extension = os.path.splitext(value)[1]
+    if extension not in [".yaml", ".yml", ".json"]:
+        raise ValueError(f"Environment variable {env_name} must be a YAML/JSON file")
+    try:
+        with open(value, encoding="utf-8") as fd:
+            if extension == ".json":
+                return json.load(fd)
+            return yaml.load(fd, Loader=yaml.CLoader)
+    except Exception as e:
+        raise ValueError(f"Environment variable {env_name}") from e
+def env_str(env_name: str, default: str | EllipsisType = ...) -> str:
+    """Get a string value from an environment variable."""
+    value = os.environ.get(env_name)
+    if value is None:
+        if default is ...:
+            raise ValueError(f"Missing environment variable {env_name}")
+        return default
+    return value
+def env_bool(env_name: str, default: bool | EllipsisType = ...) -> bool:
+    """Get a boolean value from an environment variable."""
+    value = os.environ.get(env_name)
+    if value is None:
+        if default is ...:
+            raise ValueError(f"Missing environment variable {env_name}")
+        return default
+    return value.lower() not in ["", "0", "false", "no", "off"]
+def env_int(env_name: str, default: int | EllipsisType = ...) -> int:
+    """Get an int value from an environment variable."""
+    value = os.environ.get(env_name)
+    if value is None:
+        if default is ...:
+            raise ValueError(f"Missing environment variable {env_name}")
+        return default
+    return int(value)
+def env_float(env_name: str, default: float | EllipsisType = ...) -> float:
+    """Get a boolean value from an environment variable."""
+    value = os.environ.get(env_name)
+    if value is None:
+        if default is ...:
+            raise ValueError(f"Missing environment variable {env_name}")
+        return default
+    return float(value)

notebooks/couchdb.http ADDED Viewed

	@@ -0,0 +1,39 @@

+@url = {{$dotenv COUCHDB_URL}}
+@user = {{$dotenv COUCHDB_USER}}
+@password = {{$dotenv COUCHDB_PASSWORD}}
+@db = {{$dotenv COUCHDB_DB}}
+@origin = 0
+@origin = 316723-g1AAAACReJzLYWBgYMpgTmHgzcvPy09JdcjLz8gvLskBCScyJNX___8_K4M5iYEpWyEXKMaekpyabGBuiK4ehwl5LECSoQFI_YcblCUMNsgsMS010dQSXVsWAD8FLCA
+@batchSize = 10
+@includeDocs = true
+###
+# @name session
+POST {{url}}/_session
+Content-Type: application/json
+{
+    "name": "{{user}}",
+    "password": "{{password}}"
+}
+###
+# @name changes
+GET {{url}}/{{db}}/_changes
+    ?limit={{batchSize}}
+    &include_docs={{includeDocs}}
+    &since={{origin}}
+###
+# @name changes
+GET {{url}}/{{db}}/_changes
+    ?limit={{batchSize}}
+    &include_docs={{includeDocs}}
+    &since={{changes.response.body.last_seq}}
+###
+# @name broken
+GET {{url}}/{{db}}/_changes
+    ?limit=10
+    &include_docs=true
+    &since=316731-g1AAAACReJzLYWBgYMpgTmHgzcvPy09JdcjLz8gvLskBCScyJNX___8_K4M5iYEpWyEXKMaekpyabGBuiK4ehwl5LECSoQFI_YcblCUNNsgsMS010dQSXVsWAD_1LCg

notebooks/models.ipynb ADDED Viewed

	@@ -0,0 +1,44 @@

+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pandas as pd\n",
+    "from fastembed import TextEmbedding\n",
+    "\n",
+    "# pd.set_option('display.max_colwidth', 100)\n",
+    "supported_models = (\n",
+    "    pd.DataFrame(TextEmbedding.list_supported_models())\n",
+    "    .sort_values(\"size_in_GB\")\n",
+    "    .drop(columns=[\"sources\", \"model_file\", \"additional_files\"])\n",
+    "    .reset_index(drop=True)\n",
+    ")\n",
+    "supported_models.style.set_properties(**{\"width\": \"300px\"}, subset=[\"description\"])"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "article-embedding-p7H3l83p-py3.12",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.7"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

notebooks/query.ipynb ADDED Viewed

	@@ -0,0 +1,168 @@

+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import warnings\n",
+    "\n",
+    "from qdrant_client import QdrantClient\n",
+    "\n",
+    "from article_embedding.embed import SentenceTransformerModel\n",
+    "\n",
+    "warnings.simplefilter(action=\"ignore\", category=FutureWarning)\n",
+    "embedding_model = SentenceTransformerModel(\"dunzhang/stella_en_400M_v5\")\n",
+    "qdrant = QdrantClient()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from qdrant_client import QdrantClient\n",
+    "\n",
+    "from article_embedding.embed import SentenceTransformerModel\n",
+    "\n",
+    "warnings.simplefilter(action=\"ignore\", category=FutureWarning)\n",
+    "embedding_model = SentenceTransformerModel(\"dunzhang/stella_en_400M_v5\")\n",
+    "qdrant = QdrantClient()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "vector = embedding_model.embed(\n",
+    "    [\n",
+    "        \"\"\"\n",
+    "The COVID-19 pandemic is the worst public health crisis since the 1918 influenza pandemic over a century ago,\n",
+    "when the science of virology was still in its infancy. By mid-July 2022, the official death toll since the\n",
+    "start of the pandemic stood at over 1 million in the US and nearly 6.4 million globally. Estimates of excess\n",
+    "deaths place the number of people killed directly or indirectly by the pandemic at over 1.2 million in the US\n",
+    "and 21.5 million internationally. In just two-and-a-half years, total deaths from the pandemic have reached the\n",
+    "estimated 20 million military and civilian deaths during the four years of World War I (1914-1918).\n",
+    "\"\"\"\n",
+    "    ]\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "vector = embedding_model.embed(\n",
+    "    [\n",
+    "        \"\"\"\n",
+    "Let us borrow Trotsky’s brilliant use of metaphor for the purpose of explaining the essential significance\n",
+    "of the petty-bourgeois pseudo-left’s call for a “Party of the 99 percent.” On the cruise ship America, one\n",
+    "percent of the passengers occupy state rooms that look out on the ocean. They dine with captain Trump in an\n",
+    "exclusive five-star restaurant, where they wash down their succulent meals with wine that costs $10,000 per\n",
+    "bottle. The \"next nine percent\" of the passengers, depending on what they can afford, make do with cabins\n",
+    "whose quality reflects their lower price. The cheapest of the rooms in the top ten percent category lack a\n",
+    "view out onto the ocean and have shabby rugs and uncomfortable mattresses. And, unless they are able and\n",
+    "willing to pay a substantial surcharge, the occupants of these rooms are not permitted to use the pool and\n",
+    "spa reserved for the richest passengers.\n",
+    "\"\"\"\n",
+    "    ]\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "vector = embedding_model.embed(\n",
+    "    [\n",
+    "        \"\"\"\n",
+    "Israel’s slaughter of tens if not hundreds of thousands of men, women and children, the destruction of entire\n",
+    "cities, hospitals, schools, mosques, and the imposition of mass starvation on Gaza’s 2.3 million residents—all\n",
+    "this has produced profound shock and anger towards the Netanyahu regime and its imperialist backers in the US,\n",
+    "Britain, Germany and elsewhere. It has fueled the political radicalisation of millions of workers and young people,\n",
+    "for whom Zionism is now synonymous with some of the worst atrocities committed anywhere since the Second World War.\n",
+    "\"\"\"\n",
+    "    ]\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "vector = embedding_model.embed([\"The COVID winter wave, the emergence of the Delta variant and the January 6th coup\"])[0]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "vector"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# calculate epoch seconds for 2023-01-01\n",
+    "query_filter = None\n",
+    "result = qdrant.query_points(\"wsws\", query=vector[0], query_filter=query_filter)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for point in result.points:\n",
+    "    doc = point.payload\n",
+    "    print(f\"{point.score * 100:.1f}% https://www.wsws.org/{doc[\"path\"]} - {doc[\"title\"]}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(result.model_dump_json(indent=2))"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

notebooks/stella.ipynb ADDED Viewed

	@@ -0,0 +1,67 @@

+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from sentence_transformers import SentenceTransformer\n",
+    "\n",
+    "# This model supports two prompts: \"s2p_query\" and \"s2s_query\" for sentence-to-passage and sentence-to-sentence tasks, respectively.\n",
+    "# They are defined in `config_sentence_transformers.json`\n",
+    "query_prompt_name = \"s2p_query\"\n",
+    "queries = [\n",
+    "    \"What are some ways to reduce stress?\",\n",
+    "    \"What are the benefits of drinking green tea?\",\n",
+    "]\n",
+    "# docs do not need any prompts\n",
+    "docs = [\n",
+    "    \"There are many effective ways to reduce stress. Some common techniques include deep breathing, meditation, and physical activity. Engaging in hobbies, spending time in nature, and connecting with loved ones can also help alleviate stress. Additionally, setting boundaries, practicing self-care, and learning to say no can prevent stress from building up.\",  # noqa: E501\n",
+    "    \"Green tea has been consumed for centuries and is known for its potential health benefits. It contains antioxidants that may help protect the body against damage caused by free radicals. Regular consumption of green tea has been associated with improved heart health, enhanced cognitive function, and a reduced risk of certain types of cancer. The polyphenols in green tea may also have anti-inflammatory and weight loss properties.\",  # noqa: E501\n",
+    "]\n",
+    "\n",
+    "# ！The default dimension is 1024, if you need other dimensions, please clone the model and modify `modules.json` to replace `2_Dense_1024` with another dimension, e.g. `2_Dense_256` or `2_Dense_8192` !  # noqa: E501\n",
+    "# on gpu\n",
+    "# model = SentenceTransformer(\"dunzhang/stella_en_400M_v5\", trust_remote_code=True).cuda()\n",
+    "# you can also use this model without the features of `use_memory_efficient_attention` and `unpad_inputs`. It can be worked in CPU.\n",
+    "model = SentenceTransformer(\n",
+    "    \"dunzhang/stella_en_400M_v5\",\n",
+    "    trust_remote_code=True,\n",
+    "    device=\"mps\",\n",
+    "    config_kwargs={\"use_memory_efficient_attention\": False, \"unpad_inputs\": False},\n",
+    ")\n",
+    "query_embeddings = model.encode(queries, prompt_name=query_prompt_name)\n",
+    "doc_embeddings = model.encode(docs)\n",
+    "print(query_embeddings.shape, doc_embeddings.shape)\n",
+    "# (2, 1024) (2, 1024)\n",
+    "\n",
+    "similarities = model.similarity(query_embeddings, doc_embeddings)\n",
+    "print(similarities)\n",
+    "# tensor([[0.8398, 0.2990],\n",
+    "#         [0.3282, 0.8095]])"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

poetry.lock ADDED Viewed

The diff for this file is too large to render. See raw diff

poetry.toml ADDED Viewed

	@@ -0,0 +1,2 @@


1	+ [virtualenvs]
2	+ in-project = true

pyproject.toml ADDED Viewed

	@@ -0,0 +1,86 @@

+[tool.poetry]
+name = "article-embedding"
+version = "0.1.0"
+description = ""
+authors = ["Jesus Lopez <jesus@jesusla.com>"]
+readme = "README.md"
+[tool.poetry.scripts]
+article-embedding = "article_embedding.cli:cli"
+[tool.poetry.dependencies]
+python = "~3.12"
+httpx = "^0.27.2"
+backoff = "^2.2.1"
+python-dotenv = "^1.0.1"
+tqdm = "^4.67.0"
+qdrant-client = { extras = ["fastembed"], version = "^1.12.1" }
+sentence-transformers = "^3.3.1"
+gspread = "^6.1.4"
+gradio = "^5.8.0"
+langchain-text-splitters = "^0.3.2"
+click = "^8.1.7"
+click-help-colors = "^0.9.4"
+modal = "^0.68.23"
+datasets = "^3.2.0"
+einops = "^0.8.0"
+torch = "^2.5.1"
+pandoc = "^2.4"
+openai = "^1.58.1"
+pyyaml = "^6.0.2"
+pydantic = "^2.10.4"
+[tool.poetry.group.dev.dependencies]
+pre-commit = "^4.0.1"
+ipykernel = "^6.29.5"
+ruff = "^0.7.4"
+nbstripout = "^0.8.1"
+[tool.poetry.group.typing.dependencies]
+mypy = "^1.13.0"
+types-tqdm = "^4.66.0.20240417"
+types-pyyaml = "^6.0.12.20241221"
+[tool.poetry.group.jupyter.dependencies]
+pandas = "^2.2.3"
+ipywidgets = "^8.1.5"
+jinja2 = "^3.1.4"
+[tool.poetry.group.ci.dependencies]
+huggingface-hub = { extras = ["cli"], version = "^0.27.1" }
+gradio = "^5.12.0"
+[tool.ruff]
+line-length = 132
+[tool.ruff.lint]
+select = [
+    # pycodestyle
+    "E",
+    # Pyflakes
+    "F",
+    # pyupgrade
+    "UP",
+    # flake8-bugbear
+    "B",
+    # flake8-simplify
+    "SIM",
+    # isort
+    "I",
+]
+[tool.mypy]
+strict = true
+[[tool.mypy.overrides]]
+module = ["gradio", "pandoc"]
+ignore_missing_imports = true
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"

requirements.txt ADDED Viewed

The diff for this file is too large to render. See raw diff

tests/__init__.py ADDED Viewed

File without changes