readme, makefile, service files

Dockerfile

@@ -0,0 +1,19 @@
+FROM python:3.11-slim
+
+WORKDIR /app
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    build-essential \
+    && rm -rf /var/lib/apt/lists/*
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY . .
+
+RUN useradd -m -u 1000 user
+USER user
+
+EXPOSE 7860
+
+CMD ["sh", "-c", "uvicorn main:app --host 0.0.0.0 --port 8000 & streamlit run app.py --server.port 7860 --server.address 0.0.0.0 --server.headless true"]

Makefile

@@ -0,0 +1,134 @@
+install:
+	pip install -r requirements.txt
+
+install-runpod:
+	pip install --ignore-installed blinker && pip install -r requirements.txt
+
+download_tedseg:
+	python -m src.datasets.download_tedseg
+
+download_pubmed:
+	python -m src.datasets.download_pubmed --k 100
+
+download_wikipedia:
+	python -m src.datasets.download_wikipedia --k 100
+
+preprocess_all: preprocess_pubmed preprocess_wikipedia preprocess_gutenberg preprocess_ted
+
+preprocess_pubmed:
+	python -m src.datasets.preprocess pubmed --input data/pubmed/raw/pubmed_raw.jsonl --output data/pubmed/pubmed_data.jsonl
+
+preprocess_wikipedia:
+	python -m src.datasets.preprocess wikipedia --input data/wikipedia/raw/wikipedia_raw.jsonl --output data/wikipedia/wikipedia_data.jsonl
+
+preprocess_gutenberg:
+	python -m src.datasets.preprocess gutenberg --raw_dir data/gutenberg/raw --output data/gutenberg/gutenberg_data.jsonl
+
+preprocess_ted:
+	python -m src.datasets.preprocess ted --input_dir data/tedseg/raw --output_dir data/tedseg
+
+create_recipes:
+	python -m src.datasets.create_recipes_dataset
+
+sentence_split_all: sentence_split_gutenberg sentence_split_pubmed sentence_split_wikipedia sentence_split_recipes
+
+sentence_split_gutenberg:
+	python -m src.datasets.sentence_splitter gutenberg --input data/gutenberg/gutenberg_data.jsonl --output data/gutenberg/gutenberg_sentences.jsonl --device cuda
+
+sentence_split_ted:
+	python -m src.datasets.sentence_splitter ted --input_dir data/tedseg/raw --output_dir data/tedseg --device cuda
+
+sentence_split_pubmed:
+	python -m src.datasets.sentence_splitter pubmed --input data/pubmed/pubmed_data.jsonl --output data/pubmed/pubmed_sentences.jsonl --device cuda
+
+sentence_split_wikipedia:
+	python -m src.datasets.sentence_splitter wikipedia --input data/wikipedia/wikipedia_data.jsonl --output data/wikipedia/wikipedia_sentences.jsonl --device cuda
+
+sentence_split_recipes:
+	python -m src.datasets.sentence_splitter recipes --input data/recipes/recipes_data.jsonl --output data/recipes/recipes_sentences.jsonl --device cuda
+
+build_recipes_pairs:
+	python -m src.datasets.build_pairs recipes
+
+train_distilbert:
+	python -m src.models.train --model distilbert --out checkpoints/distilbert
+
+train_bert:
+	python -m src.models.train --model bert --out checkpoints/bert
+
+train_deberta:
+	python -m src.models.train --model deberta --out checkpoints/deberta
+
+train_all: train_distilbert train_bert train_deberta
+
+export_distilbert:
+	python -m src.models.export_and_download --checkpoint checkpoints/distilbert/best
+
+export_bert:
+	python -m src.models.export_and_download --checkpoint checkpoints/bert/best
+
+export_deberta:
+	python -m src.models.export_and_download --checkpoint checkpoints/deberta/best
+
+export_all: export_distilbert export_bert export_deberta
+
+upload_distilbert:
+	python -m src.models.export_and_download --checkpoint checkpoints/distilbert/best --repo slamos/bc-models-distilbert
+
+upload_bert:
+	python -m src.models.export_and_download --checkpoint checkpoints/bert/best --repo slamos/bc-models-bert
+
+upload_deberta:
+	python -m src.models.export_and_download --checkpoint checkpoints/deberta/best --repo slamos/bc-models-deberta
+
+upload_all: upload_distilbert upload_bert upload_deberta
+
+download_distilbert:
+	python -m src.models.export_and_download --download distilbert
+
+download_bert:
+	python -m src.models.export_and_download --download bert
+
+download_deberta:
+	python -m src.models.export_and_download --download deberta
+
+download_all: download_distilbert download_bert download_deberta
+
+inference_distilbert:
+	python -m src.models.inference --model distilbert
+
+inference_bert:
+	python -m src.models.inference --model bert
+
+inference_deberta:
+	python -m src.models.inference --model deberta
+
+inference_local_distilbert:
+	python -m src.models.inference --model distilbert --local
+
+inference_local_bert:
+	python -m src.models.inference --model bert --local
+
+inference_local_deberta:
+	python -m src.models.inference --model deberta --local
+
+test-api:
+	python -m pytest tests/test_api.py -v
+
+test-dataset:
+	python -m pytest tests/test_dataset.py -v
+
+test-pipelines:
+	python -m pytest tests/test_pipelines.py -v
+
+test-all:
+	python -m pytest tests/ -v
+
+clean:
+	rm -rf checkpoints/distilbert checkpoints/bert checkpoints/deberta checkpoints/plots
+
+run-be:
+	uvicorn main:app --host 0.0.0.0 --port 8000
+
+run-fe:
+	streamlit run app.py

README.md

@@ -1 +0,0 @@
-# bottlecap-ml-dev-test

main.py

@@ -0,0 +1,62 @@
+import logging
+from contextlib import asynccontextmanager
+
+import torch
+from dotenv import load_dotenv
+from fastapi import FastAPI
+
+from src.api import fix_newlines, fix_newlines_all_models, health
+from src.models.export_and_download import download_all_models
+from src.pipelines.all_models_pipeline import AllModelsPipeline
+from src.pipelines.one_model_pipeline import OneModelPipeline
+from src.pipelines.sat_loader import load_sat
+
+load_dotenv()
+log = logging.getLogger(__name__)
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    log.info(f"Loading SAT-3L-SM on {device} ...")
+    sat = load_sat(device=device)
+
+    downloaded = download_all_models()
+
+    # OneModelPipeline — uses bert
+    bert_dir = downloaded.get("bert")
+    if bert_dir and (bert_dir / "model.onnx").exists():
+        log.info("Loading OneModelPipeline (bert)")
+        app.state.one_model_pipeline = OneModelPipeline(
+            onnx_path=bert_dir / "model.onnx",
+            tokenizer_path=bert_dir,
+            sat_model=sat,
+            model_name="bert",
+        )
+    else:
+        log.warning("bert not available — OneModelPipeline disabled")
+        app.state.one_model_pipeline = None
+
+    # AllModelsPipeline — uses all downloaded models
+    all_pipeline = AllModelsPipeline(sat_model=sat)
+    for name, local_dir in downloaded.items():
+        onnx_path = local_dir / "model.onnx"
+        if onnx_path.exists():
+            log.info(f"Adding {name} to AllModelsPipeline")
+            all_pipeline.add_model(name, onnx_path, local_dir)
+        else:
+            log.warning(f"Skipping {name}: model.onnx not found in {local_dir}")
+    app.state.all_models_pipeline = all_pipeline
+
+    yield
+
+
+app = FastAPI(
+    title="Newline Fixer API",
+    description="ML service for fixing newline placement in English text",
+    lifespan=lifespan,
+)
+
+app.include_router(health.router)
+app.include_router(fix_newlines.router)
+app.include_router(fix_newlines_all_models.router)

pages/nav.py

@@ -0,0 +1,32 @@
+"""Shared navigation stepper for all pages."""
+import streamlit as st
+
+PAGES = [
+    ("Home", "pages/home.py"),
+    ("Config", "pages/config.py"),
+    ("Result", "pages/result.py"),
+]
+
+AVAILABLE_MODELS = ["rule_based", "distilbert", "bert", "deberta"]
+
+
+def show_stepper(current: str):
+    """Render a horizontal step indicator. `current` is the page title."""
+    cols = st.columns(len(PAGES))
+    for i, (title, path) in enumerate(PAGES):
+        with cols[i]:
+            if title == current:
+                st.markdown(
+                    f"<div style='text-align:center; padding:8px 0; "
+                    f"border-bottom:3px solid #ff4b4b; font-weight:700;'>"
+                    f"{i + 1}. {title}</div>",
+                    unsafe_allow_html=True,
+                )
+            else:
+                st.markdown(
+                    f"<div style='text-align:center; padding:8px 0; "
+                    f"border-bottom:3px solid #444; color:#888;'>"
+                    f"{i + 1}. {title}</div>",
+                    unsafe_allow_html=True,
+                )
+    st.write("")

pytest.ini

@@ -0,0 +1,3 @@
+[pytest]
+filterwarnings =
+    ignore:split_on_input_newlines=False:UserWarning

requirements.txt

@@ -0,0 +1,22 @@
+fastapi>=0.121.0
+uvicorn>=0.34.0
+pydantic>=2.12.4
+torch>=2.6.0
+torchvision>=0.21.0
+transformers>=4.50.3
+datasets>=2.21.0
+tqdm>=4.67.3
+pandas>=1.4.0
+requests>=2.32.5
+python-dotenv>=1.2.1
+wtpsplit>=2.2.0
+sentencepiece>=0.2.0
+onnx>=1.16.0
+onnxruntime>=1.18.0
+onnxscript>=0.1.0
+huggingface_hub>=0.23.0
+matplotlib
+accelerate>=0.26.0
+wandb>=0.17.0
+streamlit>=1.35.0
+pytest>=8.0.0

src/datasets/preprocess.py

@@ -0,0 +1,182 @@
+import argparse
+import json
+import re
+from pathlib import Path
+
+
+def preprocess_gutenberg(raw_dir: str, output_path: str):
+    """
+    Read all files from raw_dir (Gutenberg books), split each into paragraphs,
+    and write a single JSONL file where each line is one book:
+
+        {"file_name": "3046.ref", "text": ["paragraph 1", "paragraph 2", ...]}
+
+    Paragraph boundaries are determined by:
+      - double newlines (\\n\\n)
+      - lines starting with capitalised text, a number, or ==========
+    """
+    raw = Path(raw_dir)
+    out = Path(output_path)
+    out.parent.mkdir(parents=True, exist_ok=True)
+
+    n_books = 0
+    with open(out, "w", encoding="utf-8") as fout:
+        for filepath in sorted(raw.iterdir()):
+            if filepath.is_dir():
+                continue
+            text = filepath.read_text(encoding="utf-8")
+
+            # split on ========== into sections
+            sections = re.split(r"\n*==========\n*", text)
+
+            paragraphs = []
+            for section in sections:
+                section = section.strip()
+                if not section:
+                    continue
+                # first line is the heading, rest is body text
+                first_newline = section.find("\n")
+                if first_newline == -1:
+                    # section is just a heading with no body
+                    paragraphs.append(section)
+                else:
+                    heading = section[:first_newline].strip()
+                    body = section[first_newline + 1:].strip()
+                    if heading:
+                        paragraphs.append(heading)
+                    if body:
+                        paragraphs.append(body)
+
+            if not paragraphs:
+                continue
+
+            record = {"file_name": filepath.name, "text": paragraphs}
+            fout.write(json.dumps(record, ensure_ascii=False) + "\n")
+            n_books += 1
+
+    print(f"Wrote {n_books:,} books to {out}")
+
+
+def preprocess_ted(input_dir: str, output_dir: str, max_paragraphs: int = 60):
+    """
+    Process all TED split files (train.jsonl, val.jsonl, test.jsonl) in input_dir,
+    drop samples with more than max_paragraphs paragraphs, and write each as
+    {name}_preprocessed.jsonl into output_dir.
+    """
+    inp = Path(input_dir)
+    out = Path(output_dir)
+    out.mkdir(parents=True, exist_ok=True)
+
+    for split_file in sorted(inp.glob("*.jsonl")):
+        out_file = out / f"{split_file.stem}_preprocessed.jsonl"
+        n_kept = 0
+        n_dropped = 0
+        with open(split_file, encoding="utf-8") as fin, open(out_file, "w", encoding="utf-8") as fout:
+            for line in fin:
+                line = line.strip()
+                if not line:
+                    continue
+                obj = json.loads(line)
+                paragraphs = obj.get("text", [])
+                if len(paragraphs) > max_paragraphs:
+                    n_dropped += 1
+                    continue
+                fout.write(json.dumps(obj, ensure_ascii=False) + "\n")
+                n_kept += 1
+        print(f"{split_file.name}: kept {n_kept:,}, dropped {n_dropped:,} with >{max_paragraphs} paragraphs → {out_file}")
+
+
+def preprocess_jsonl_paragraphs(input_path: str, output_path: str, text_field: str = "text"):
+    """
+    Read a raw JSONL file, split the text field on double newlines into
+    a list of paragraphs, and write the result as a new JSONL file.
+
+    Output format per line:
+        {"document_idx": 0, "text": ["paragraph 1", "paragraph 2", ...], ...}
+    """
+    inp = Path(input_path)
+    out = Path(output_path)
+    out.parent.mkdir(parents=True, exist_ok=True)
+
+    n_docs = 0
+    with open(inp, encoding="utf-8") as fin, open(out, "w", encoding="utf-8") as fout:
+        for line in fin:
+            line = line.strip()
+            if not line:
+                continue
+            obj = json.loads(line)
+            raw_text = obj.get(text_field, "")
+            paragraphs = [p.strip() for p in raw_text.split("\n\n") if p.strip()]
+            obj[text_field] = paragraphs
+            fout.write(json.dumps(obj, ensure_ascii=False) + "\n")
+            n_docs += 1
+
+    print(f"Wrote {n_docs:,} documents to {out}")
+
+
+def _strip_yaml_frontmatter(text: str) -> str:
+    """Remove YAML frontmatter (--- ... ---) from the beginning of text."""
+    if text.startswith("---"):
+        end = text.find("\n---", 3)
+        if end != -1:
+            text = text[end + 4:].lstrip("\n")
+    return text
+
+
+def preprocess_pubmed(input_path: str, output_path: str):
+    inp = Path(input_path)
+    out = Path(output_path)
+    out.parent.mkdir(parents=True, exist_ok=True)
+
+    n_docs = 0
+    with open(inp, encoding="utf-8") as fin, open(out, "w", encoding="utf-8") as fout:
+        for line in fin:
+            line = line.strip()
+            if not line:
+                continue
+            obj = json.loads(line)
+            raw_text = _strip_yaml_frontmatter(obj.get("text", ""))
+            paragraphs = [p.strip() for p in raw_text.split("\n\n") if p.strip()]
+            obj["text"] = paragraphs
+            fout.write(json.dumps(obj, ensure_ascii=False) + "\n")
+            n_docs += 1
+
+    print(f"Wrote {n_docs:,} documents to {out}")
+
+
+def preprocess_wikipedia(input_path: str, output_path: str):
+    preprocess_jsonl_paragraphs(input_path, output_path, text_field="text")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Preprocess raw data into JSONL.")
+    sub = parser.add_subparsers(dest="method", required=True)
+
+    gut = sub.add_parser("gutenberg", help="Convert raw Gutenberg books to JSONL")
+    gut.add_argument("--raw_dir", required=True, help="Directory with raw Gutenberg text files")
+    gut.add_argument("--output", required=True, help="Output JSONL path")
+
+    ted = sub.add_parser("ted", help="Filter TED talks JSONL by paragraph count")
+    ted.add_argument("--input_dir", required=True, help="Directory with TED split JSONL files")
+    ted.add_argument("--output_dir", required=True, help="Output directory for preprocessed files")
+    ted.add_argument("--max_paragraphs", type=int, default=60,
+                     help="Drop talks with more than this many paragraphs (default: 60)")
+
+    pub = sub.add_parser("pubmed", help="Split PubMed text into paragraphs")
+    pub.add_argument("--input", required=True, help="Input raw JSONL path")
+    pub.add_argument("--output", required=True, help="Output JSONL path")
+
+    wiki = sub.add_parser("wikipedia", help="Split Wikipedia text into paragraphs")
+    wiki.add_argument("--input", required=True, help="Input raw JSONL path")
+    wiki.add_argument("--output", required=True, help="Output JSONL path")
+
+    args = parser.parse_args()
+
+    if args.method == "gutenberg":
+        preprocess_gutenberg(args.raw_dir, args.output)
+    elif args.method == "ted":
+        preprocess_ted(args.input_dir, args.output_dir, args.max_paragraphs)
+    elif args.method == "pubmed":
+        preprocess_pubmed(args.input, args.output)
+    elif args.method == "wikipedia":
+        preprocess_wikipedia(args.input, args.output)

src/logger.py

@@ -0,0 +1,47 @@
+"""Betterstack logging for API endpoint calls."""
+
+import os
+import time
+from datetime import datetime, timezone
+
+import requests
+
+
+BETTERSTACK_URL = "https://in.logs.betterstack.com"
+
+
+def log_to_betterstack(
+    endpoint_name: str,
+    original_text: str,
+    formatted_text: str | dict,
+    model_name: str,
+    time_elapsed: float,
+) -> None:
+    token = os.environ.get("BETTERSTACK_SOURCE_TOKEN")
+    if not token:
+        return
+
+    payload = {
+        "dt": datetime.now(timezone.utc).isoformat(),
+        "message": f"Endpoint called: {endpoint_name}",
+        "endpoint_name": endpoint_name,
+        "model_name": model_name,
+        "original_text": original_text,
+        "formatted_text": formatted_text,
+        "time_elapsed": round(time_elapsed, 4),
+        "request_date": datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S"),
+        "level": "info"
+    }
+
+    try:
+        requests.post(
+            BETTERSTACK_URL,
+            json=payload,
+            headers={
+                "Authorization": f"Bearer {token}",
+                "Content-Type": "application/json",
+            },
+            timeout=5,
+        )
+    except Exception:
+        pass

src/schemas/requests.py

@@ -0,0 +1,9 @@
+from pydantic import BaseModel, Field
+
+
+class FixNewlinesRequest(BaseModel):
+    text: str = Field(..., description="Input text with broken newline placement")
+
+
+class FixNewlinesAllModelsRequest(BaseModel):
+    text: str = Field(..., description="Input text with broken newline placement")

src/schemas/responses.py

@@ -0,0 +1,22 @@
+from pydantic import BaseModel, Field
+
+
+class FixNewlinesResponse(BaseModel):
+    fixed_text: str = Field(..., description="Text with corrected newline placement")
+    model_used: str = Field(..., description="Name of the model that produced the result")
+
+
+class ModelResult(BaseModel):
+    model_name: str
+    fixed_text: str
+
+
+class FixNewlinesAllModelsResponse(BaseModel):
+    results: list[ModelResult] = Field(
+        ..., description="Results from all available models"
+    )
+
+
+class HealthResponse(BaseModel):
+    status: str
+    available_models: list[str]

task_description.md

@@ -0,0 +1,80 @@
+# Applied ML Engineer Challenge
+
+![logo](img/logo.png)
+
+Hey there! Are you interested in LLMs? Do you like building real products, experimenting with neural networks, implementing different ideas and testing them out? Would you like to do that for a living? Then you're in the right place!
+This is an official test for people interested in joining [BottleCapAI](https://www.bottlecapai.com).
+
+---
+
+## Objective
+
+Design and implement a Machine Learning service capable of fixing newline placement in English natural language text. Develop a model,
+and a service with a HTTP API. You can choose any architecture and method of obtaining the model, and any Python framework for the HTTP API. 
+It is not necessary to aim for state-of-the-art performance, but the service should efficiently generate reasonable answers.
+
+Examples:  
+```
+3.2.3 Applications of Attention
+ in our Model The Transformer uses multi-head attention in three different ways: • In "encoder-decoder attention" layers,
+ the que
+ries come from the previous decoder layer.
+• In "encoder-only attention" layers, 
+the que
+ries come from the previous decoder layer.[...]
+``` 
+-> 
+```
+3.2.3 Applications of Attention in our Model
+
+The Transformer uses multi-head attention in three different ways:
+• In "encoder-decoder attention" layers, the queries come from the previous decoder layer.
+[...]
+```
+
+For your solution, it is required to:
+- Implement an API with at least one endpoint that accepts a text input and returns a text with fixed newlines.
+- Provide a Dockerfile with the environment to run the service.
+- Write tests for the service.
+- Compute appropriate metrics for evaluating the performance of the model.
+- Provide a `report.md` with instructions on how to run the service, explaining your approach, decisions taken, and reporting the results.
+- [optional] Ideally, we would encourage you to provide us with a link to the service deployed, for instance on [Huggingface Spaces](https://huggingface.co/spaces), with a minimal UI
+for testing the model (can be in [Streamlit](https://github.com/streamlit/streamlit) or basic HTML + js). Provide the space link in your report.
+
+In case you need access to GPUs for developing your model, we recommend using free online solutions, such as [Google Colab](https://colab.research.google.com/),
+[Modal notebooks](https://modal.com/products/notebooks), or [Kaggle kernels](https://www.kaggle.com/kernels). 
+It is fine to focus on smaller models in case of any hardware-related difficulties, as we are not expecting SOTA performance.
+
+---
+
+## What's the point?
+
+We are interested in your drive, interest in ML and ability to create useful working products with it, adapting models to certain tasks, 
+skill in developing, evaluating, monitoring, and deploying production ready ML services, and general programming skills.
+
+---
+
+## Submission
+
+To submit your results, run:
+```bash
+git bundle create <first name>-<last name>.bundle --all
+```
+Then send us your .bundle file to hey(at)bottlecapai.com with subject in format: \<first name\>-\<last name\>-applied-ml-test\>.
+
+At this moment, we are interested mainly in candidates willing to relocate to Prague and authorized to work in the EU. (If you are an exceptional fit, we are happy to discuss possible support options).
+
+---
+## 📌 About BottleCapAI
+
+At **BottleCapAI**, we’re making large language models **radically more efficient** — aiming for **100× improvements** over today’s approaches. 🚀  
+
+### 👥 Founders
+- Tomas Mikolov – creator of *word2vec*, pioneer of neural language models.  
+- Jaroslav Beck – co-founder of *Beat Games* (*Beat Saber*, 10M+ copies sold, acquired by Meta).  
+- David Herel – creator of Thinking Tokens, co-founder of an AI trading startup, and Amazon Alexa Prize finalist.
+
+### 🌍 Our vision
+Training frontier LLMs costs **tens of millions** today. Our new algorithms already cut that by **~50%** — and we’re just getting started. We’re building a European hub to push AI forward through **algorithms, not brute force**.  
+ 
+📧 **hey(at)bottlecapai.com** · 🌐 [bottlecapai.com](https://www.bottlecapai.com)