feat: free alternatives to all paid APIs — local fine-tune, OCR, STT, autorouter
Tools (zero API cost): - local_finetune.py: Unsloth QLoRA on RTX 4090 → GGUF → Ollama - ocr_pipeline.py: marker/surya/PyPDF2 for datasheet extraction - stt_pipeline.py: whisper.cpp/vosk for meeting transcription - freerouting_bridge.py: open source autorouter for KiCad PCBs Plans closed with free alternatives: - Plan 24: 15 tasks (fine-tune, OCR, STT, RAG, deploy) - Plan 23/23v2: 3 tasks (fine-tune, benchmark) - Plan 25: 4 tasks (SPICE, eval, Quilter, PCBDesigner) - Plan 26: 1 task (Hypnoled validation) - Plan 27: 1 task (Jetson → Docker NVIDIA) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
L'électron rare
@@ -0,0 +1,359 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Freerouting Bridge — FREE alternative to paid autorouting services (T-HP-033/034).
|
||||
|
||||
Bridges KiCad DSN export -> Freerouting (open source Java autorouter) -> KiCad SES import.
|
||||
|
||||
Freerouting: https://github.com/freerouting/freerouting
|
||||
|
||||
Usage:
|
||||
# Route a board (auto-downloads Freerouting JAR if needed)
|
||||
python3 freerouting_bridge.py route --input board.dsn --output board.ses
|
||||
|
||||
# Just download / update Freerouting
|
||||
python3 freerouting_bridge.py download
|
||||
|
||||
# Verify DSN file before routing
|
||||
python3 freerouting_bridge.py check --input board.dsn
|
||||
|
||||
# Specify custom JAR location
|
||||
python3 freerouting_bridge.py route --input board.dsn --jar /path/to/freerouting.jar
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import platform
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Constants
|
||||
# ---------------------------------------------------------------------------
|
||||
FREEROUTING_GITHUB = "freerouting/freerouting"
|
||||
FREEROUTING_RELEASE_API = f"https://api.github.com/repos/{FREEROUTING_GITHUB}/releases/latest"
|
||||
DEFAULT_JAR_DIR = Path.home() / ".local" / "share" / "freerouting"
|
||||
DEFAULT_JAR_PATH = DEFAULT_JAR_DIR / "freerouting.jar"
|
||||
|
||||
# Minimum Java version
|
||||
MIN_JAVA_VERSION = 17
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Java detection
|
||||
# ---------------------------------------------------------------------------
|
||||
def find_java() -> Optional[str]:
|
||||
"""Find a suitable Java runtime."""
|
||||
# Check JAVA_HOME first
|
||||
java_home = os.environ.get("JAVA_HOME")
|
||||
if java_home:
|
||||
java_bin = Path(java_home) / "bin" / "java"
|
||||
if java_bin.exists():
|
||||
return str(java_bin)
|
||||
|
||||
# Check PATH
|
||||
java = shutil.which("java")
|
||||
if java:
|
||||
return java
|
||||
|
||||
# macOS: check common Homebrew / system locations
|
||||
if platform.system() == "Darwin":
|
||||
candidates = [
|
||||
"/opt/homebrew/bin/java",
|
||||
"/usr/local/bin/java",
|
||||
"/usr/bin/java",
|
||||
]
|
||||
for c in candidates:
|
||||
if Path(c).exists():
|
||||
return c
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def check_java_version(java_bin: str) -> int:
|
||||
"""Return Java major version number, or 0 on failure."""
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[java_bin, "-version"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=10,
|
||||
)
|
||||
output = result.stderr + result.stdout
|
||||
match = re.search(r'"(\d+)', output)
|
||||
if match:
|
||||
return int(match.group(1))
|
||||
except Exception:
|
||||
pass
|
||||
return 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Freerouting JAR management
|
||||
# ---------------------------------------------------------------------------
|
||||
def download_freerouting(dest: Path = DEFAULT_JAR_PATH, force: bool = False) -> Path:
|
||||
"""Download the latest Freerouting JAR from GitHub releases."""
|
||||
dest.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if dest.exists() and not force:
|
||||
print(f" Freerouting already present: {dest}")
|
||||
print(f" Use --force to re-download")
|
||||
return dest
|
||||
|
||||
print(f" Fetching latest release from {FREEROUTING_GITHUB} ...")
|
||||
|
||||
try:
|
||||
req = urllib.request.Request(
|
||||
FREEROUTING_RELEASE_API,
|
||||
headers={"Accept": "application/vnd.github.v3+json", "User-Agent": "kill-life-bridge"},
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=30) as resp:
|
||||
release = json.loads(resp.read())
|
||||
except Exception as exc:
|
||||
sys.exit(f"ERROR: failed to fetch release info: {exc}")
|
||||
|
||||
# Find the JAR asset
|
||||
jar_asset = None
|
||||
for asset in release.get("assets", []):
|
||||
name = asset["name"].lower()
|
||||
if name.endswith(".jar") and "freerouting" in name:
|
||||
jar_asset = asset
|
||||
break
|
||||
|
||||
if not jar_asset:
|
||||
# Some releases use the executable JAR without "freerouting" in the name
|
||||
for asset in release.get("assets", []):
|
||||
if asset["name"].lower().endswith(".jar"):
|
||||
jar_asset = asset
|
||||
break
|
||||
|
||||
if not jar_asset:
|
||||
sys.exit(
|
||||
f"ERROR: no JAR found in release {release.get('tag_name', '?')}.\n"
|
||||
f" Download manually from https://github.com/{FREEROUTING_GITHUB}/releases"
|
||||
)
|
||||
|
||||
download_url = jar_asset["browser_download_url"]
|
||||
size_mb = jar_asset.get("size", 0) / (1024 * 1024)
|
||||
print(f" Downloading {jar_asset['name']} ({size_mb:.1f} MB) ...")
|
||||
|
||||
try:
|
||||
urllib.request.urlretrieve(download_url, str(dest))
|
||||
except Exception as exc:
|
||||
sys.exit(f"ERROR: download failed: {exc}")
|
||||
|
||||
print(f" -> {dest}")
|
||||
return dest
|
||||
|
||||
|
||||
def find_jar(jar_path: Optional[str] = None) -> Path:
|
||||
"""Locate the Freerouting JAR."""
|
||||
if jar_path:
|
||||
p = Path(jar_path)
|
||||
if p.exists():
|
||||
return p
|
||||
sys.exit(f"ERROR: JAR not found at {jar_path}")
|
||||
|
||||
# Check environment variable
|
||||
env_jar = os.environ.get("FREEROUTING_JAR")
|
||||
if env_jar and Path(env_jar).exists():
|
||||
return Path(env_jar)
|
||||
|
||||
# Check default location
|
||||
if DEFAULT_JAR_PATH.exists():
|
||||
return DEFAULT_JAR_PATH
|
||||
|
||||
# Auto-download
|
||||
print(" Freerouting JAR not found, downloading ...")
|
||||
return download_freerouting()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# DSN validation
|
||||
# ---------------------------------------------------------------------------
|
||||
def check_dsn(dsn_path: Path) -> Dict[str, Any]:
|
||||
"""Basic validation of a KiCad DSN file."""
|
||||
if not dsn_path.exists():
|
||||
sys.exit(f"ERROR: file not found: {dsn_path}")
|
||||
|
||||
text = dsn_path.read_text(encoding="utf-8", errors="replace")
|
||||
|
||||
info: Dict[str, Any] = {
|
||||
"file": str(dsn_path),
|
||||
"size_bytes": dsn_path.stat().st_size,
|
||||
"valid": False,
|
||||
}
|
||||
|
||||
# Check for DSN header
|
||||
if not text.strip().startswith("(pcb"):
|
||||
info["error"] = "File does not start with (pcb — not a valid DSN export"
|
||||
return info
|
||||
|
||||
# Count components and nets
|
||||
components = re.findall(r"\(component\s", text)
|
||||
nets = re.findall(r"\(net\s", text)
|
||||
wires = re.findall(r"\(wire\s", text)
|
||||
vias = re.findall(r"\(via\s", text)
|
||||
|
||||
info.update({
|
||||
"valid": True,
|
||||
"components": len(components),
|
||||
"nets": len(nets),
|
||||
"existing_wires": len(wires),
|
||||
"existing_vias": len(vias),
|
||||
})
|
||||
|
||||
# Check paren balance
|
||||
opens = text.count("(")
|
||||
closes = text.count(")")
|
||||
if opens != closes:
|
||||
info["warning"] = f"Unbalanced parentheses: {opens} open vs {closes} close"
|
||||
|
||||
return info
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Routing
|
||||
# ---------------------------------------------------------------------------
|
||||
def route(
|
||||
dsn_path: Path,
|
||||
output_path: Optional[Path] = None,
|
||||
jar_path: Optional[str] = None,
|
||||
java_bin: Optional[str] = None,
|
||||
timeout_seconds: int = 600,
|
||||
extra_args: Optional[List[str]] = None,
|
||||
) -> Path:
|
||||
"""Run Freerouting on a DSN file, producing a SES file."""
|
||||
dsn_path = dsn_path.resolve()
|
||||
if not dsn_path.exists():
|
||||
sys.exit(f"ERROR: DSN file not found: {dsn_path}")
|
||||
|
||||
# Validate DSN
|
||||
info = check_dsn(dsn_path)
|
||||
if not info["valid"]:
|
||||
sys.exit(f"ERROR: invalid DSN file: {info.get('error', 'unknown')}")
|
||||
print(f" DSN: {info['components']} components, {info['nets']} nets")
|
||||
|
||||
# Find Java
|
||||
java = java_bin or find_java()
|
||||
if not java:
|
||||
sys.exit(
|
||||
"ERROR: Java not found. Install Java >= 17:\n"
|
||||
" macOS: brew install openjdk@17\n"
|
||||
" Linux: sudo apt install openjdk-17-jre-headless"
|
||||
)
|
||||
version = check_java_version(java)
|
||||
if version and version < MIN_JAVA_VERSION:
|
||||
print(f" WARN: Java {version} detected, Freerouting needs >= {MIN_JAVA_VERSION}")
|
||||
|
||||
# Find JAR
|
||||
jar = find_jar(jar_path)
|
||||
|
||||
# Output path
|
||||
if output_path is None:
|
||||
output_path = dsn_path.with_suffix(".ses")
|
||||
|
||||
# Build command
|
||||
cmd = [
|
||||
java,
|
||||
"-jar", str(jar),
|
||||
"-de", str(dsn_path), # design input
|
||||
"-do", str(output_path), # design output (SES)
|
||||
"-mp", "20", # max passes
|
||||
]
|
||||
if extra_args:
|
||||
cmd.extend(extra_args)
|
||||
|
||||
print(f" Running Freerouting (timeout={timeout_seconds}s) ...")
|
||||
print(f" {' '.join(cmd)}")
|
||||
|
||||
try:
|
||||
result = subprocess.run(
|
||||
cmd,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=timeout_seconds,
|
||||
env={**os.environ, "DISPLAY": ""}, # headless
|
||||
)
|
||||
except subprocess.TimeoutExpired:
|
||||
sys.exit(f"ERROR: Freerouting timed out after {timeout_seconds}s")
|
||||
|
||||
if result.returncode != 0:
|
||||
stderr = result.stderr.strip()
|
||||
stdout = result.stdout.strip()
|
||||
# Freerouting may still produce output even with non-zero exit
|
||||
if output_path.exists() and output_path.stat().st_size > 0:
|
||||
print(f" WARN: Freerouting exited with code {result.returncode} but produced output")
|
||||
else:
|
||||
sys.exit(
|
||||
f"ERROR: Freerouting failed (exit {result.returncode}):\n"
|
||||
f" stderr: {stderr[:500]}\n"
|
||||
f" stdout: {stdout[:500]}"
|
||||
)
|
||||
|
||||
if not output_path.exists():
|
||||
sys.exit("ERROR: Freerouting produced no SES output")
|
||||
|
||||
print(f" -> {output_path} ({output_path.stat().st_size} bytes)")
|
||||
print(f"\n Import into KiCad:")
|
||||
print(f" File -> Import -> Specctra Session (.ses)")
|
||||
return output_path
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CLI
|
||||
# ---------------------------------------------------------------------------
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Freerouting bridge for KiCad (free autorouting)",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog=__doc__,
|
||||
)
|
||||
sub = parser.add_subparsers(dest="command")
|
||||
|
||||
# route
|
||||
p_route = sub.add_parser("route", help="Route a DSN file")
|
||||
p_route.add_argument("--input", type=Path, required=True, help="KiCad DSN export")
|
||||
p_route.add_argument("--output", type=Path, help="SES output (default: same name .ses)")
|
||||
p_route.add_argument("--jar", type=str, help="Path to freerouting.jar")
|
||||
p_route.add_argument("--java", type=str, help="Path to java binary")
|
||||
p_route.add_argument("--timeout", type=int, default=600, help="Timeout in seconds")
|
||||
|
||||
# download
|
||||
p_dl = sub.add_parser("download", help="Download/update Freerouting JAR")
|
||||
p_dl.add_argument("--dest", type=Path, default=DEFAULT_JAR_PATH)
|
||||
p_dl.add_argument("--force", action="store_true")
|
||||
|
||||
# check
|
||||
p_chk = sub.add_parser("check", help="Validate a DSN file")
|
||||
p_chk.add_argument("--input", type=Path, required=True)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.command == "route":
|
||||
route(
|
||||
dsn_path=args.input,
|
||||
output_path=args.output,
|
||||
jar_path=args.jar,
|
||||
java_bin=args.java,
|
||||
timeout_seconds=args.timeout,
|
||||
)
|
||||
elif args.command == "download":
|
||||
download_freerouting(dest=args.dest, force=args.force)
|
||||
elif args.command == "check":
|
||||
info = check_dsn(args.input)
|
||||
print(json.dumps(info, indent=2))
|
||||
else:
|
||||
parser.print_help()
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,318 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
OCR Pipeline — FREE alternative to Mistral's document parsing API (T-MS-020).
|
||||
|
||||
Extracts text and structured component specs from PDF datasheets using
|
||||
local/open-source tools, with graceful fallback chain:
|
||||
|
||||
1. marker-pdf (best quality, GPU-accelerated)
|
||||
2. surya-ocr (good quality, lighter)
|
||||
3. PyPDF2 (text-layer only, no OCR)
|
||||
|
||||
Usage:
|
||||
# Single PDF
|
||||
python3 ocr_pipeline.py --pdf datasheet.pdf --output specs.json
|
||||
|
||||
# Batch directory
|
||||
python3 ocr_pipeline.py --dir datasheets/ --output-dir specs/
|
||||
|
||||
# Force a specific backend
|
||||
python3 ocr_pipeline.py --pdf datasheet.pdf --backend pypdf2 --output specs.json
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Backend: marker-pdf
|
||||
# ---------------------------------------------------------------------------
|
||||
def _ocr_marker(pdf_path: Path) -> str:
|
||||
"""Use marker-pdf to convert PDF to markdown text."""
|
||||
try:
|
||||
from marker.converters.pdf import PdfConverter
|
||||
from marker.models import create_model_dict
|
||||
except ImportError:
|
||||
raise RuntimeError("marker-pdf not installed: pip install marker-pdf")
|
||||
|
||||
models = create_model_dict()
|
||||
converter = PdfConverter(artifact_dict=models)
|
||||
rendered = converter(str(pdf_path))
|
||||
return rendered.markdown
|
||||
|
||||
|
||||
def _ocr_marker_cli(pdf_path: Path) -> str:
|
||||
"""Fallback: call marker CLI as subprocess."""
|
||||
import subprocess
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
result = subprocess.run(
|
||||
["marker_single", str(pdf_path), tmp, "--batch_multiplier", "2"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=300,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
raise RuntimeError(f"marker_single failed: {result.stderr[:500]}")
|
||||
# marker writes a .md file in the output dir
|
||||
md_files = list(Path(tmp).rglob("*.md"))
|
||||
if not md_files:
|
||||
raise RuntimeError("marker produced no output")
|
||||
return md_files[0].read_text(encoding="utf-8")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Backend: surya-ocr
|
||||
# ---------------------------------------------------------------------------
|
||||
def _ocr_surya(pdf_path: Path) -> str:
|
||||
"""Use surya for OCR."""
|
||||
try:
|
||||
from surya.ocr import run_ocr
|
||||
from surya.model.detection.model import load_model as load_det_model
|
||||
from surya.model.detection.processor import load_processor as load_det_proc
|
||||
from surya.model.recognition.model import load_model as load_rec_model
|
||||
from surya.model.recognition.processor import load_processor as load_rec_proc
|
||||
from surya.input.load import load_from_file
|
||||
except ImportError:
|
||||
raise RuntimeError("surya-ocr not installed: pip install surya-ocr")
|
||||
|
||||
det_model = load_det_model()
|
||||
det_proc = load_det_proc()
|
||||
rec_model = load_rec_model()
|
||||
rec_proc = load_rec_proc()
|
||||
|
||||
images, _ = load_from_file(str(pdf_path))
|
||||
langs = [["en"]] * len(images)
|
||||
|
||||
results = run_ocr(
|
||||
images, langs, det_model, det_proc, rec_model, rec_proc
|
||||
)
|
||||
|
||||
pages: List[str] = []
|
||||
for page_result in results:
|
||||
lines = [line.text for line in page_result.text_lines]
|
||||
pages.append("\n".join(lines))
|
||||
return "\n\n---\n\n".join(pages)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Backend: PyPDF2 (text-layer only, no OCR)
|
||||
# ---------------------------------------------------------------------------
|
||||
def _ocr_pypdf2(pdf_path: Path) -> str:
|
||||
"""Extract embedded text layer — no actual OCR."""
|
||||
try:
|
||||
from PyPDF2 import PdfReader
|
||||
except ImportError:
|
||||
try:
|
||||
from pypdf import PdfReader
|
||||
except ImportError:
|
||||
raise RuntimeError("PyPDF2/pypdf not installed: pip install pypdf")
|
||||
|
||||
reader = PdfReader(str(pdf_path))
|
||||
pages: List[str] = []
|
||||
for page in reader.pages:
|
||||
text = page.extract_text()
|
||||
if text:
|
||||
pages.append(text)
|
||||
if not pages:
|
||||
raise RuntimeError("PDF has no embedded text layer (scanned image?)")
|
||||
return "\n\n---\n\n".join(pages)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Backend dispatcher with fallback chain
|
||||
# ---------------------------------------------------------------------------
|
||||
BACKENDS = [
|
||||
("marker", _ocr_marker),
|
||||
("marker_cli", _ocr_marker_cli),
|
||||
("surya", _ocr_surya),
|
||||
("pypdf2", _ocr_pypdf2),
|
||||
]
|
||||
|
||||
|
||||
def extract_text(pdf_path: Path, backend: Optional[str] = None) -> tuple[str, str]:
|
||||
"""
|
||||
Extract text from a PDF. Returns (text, backend_used).
|
||||
Falls through the chain on failure unless a specific backend is requested.
|
||||
"""
|
||||
if backend:
|
||||
# Direct backend selection
|
||||
for name, fn in BACKENDS:
|
||||
if name == backend:
|
||||
return fn(pdf_path), name
|
||||
sys.exit(f"ERROR: unknown backend '{backend}'. Choose from: {[b[0] for b in BACKENDS]}")
|
||||
|
||||
errors: List[str] = []
|
||||
for name, fn in BACKENDS:
|
||||
try:
|
||||
text = fn(pdf_path)
|
||||
if text.strip():
|
||||
return text, name
|
||||
errors.append(f"{name}: empty output")
|
||||
except Exception as exc:
|
||||
errors.append(f"{name}: {exc}")
|
||||
|
||||
sys.exit(
|
||||
f"ERROR: all OCR backends failed for {pdf_path.name}:\n"
|
||||
+ "\n".join(f" - {e}" for e in errors)
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Structured spec extraction (regex-based, no LLM needed)
|
||||
# ---------------------------------------------------------------------------
|
||||
def extract_specs(text: str, filename: str = "") -> Dict[str, Any]:
|
||||
"""
|
||||
Pull common component specs from OCR text via regex patterns.
|
||||
Returns a JSON-serialisable dict.
|
||||
"""
|
||||
specs: Dict[str, Any] = {"source_file": filename}
|
||||
|
||||
# Voltage ratings
|
||||
voltages = re.findall(
|
||||
r"(\d+(?:\.\d+)?)\s*(?:V(?:DC|AC|RMS)?|volts?)\b", text, re.IGNORECASE
|
||||
)
|
||||
if voltages:
|
||||
specs["voltages_V"] = sorted(set(float(v) for v in voltages))
|
||||
|
||||
# Current ratings
|
||||
currents = re.findall(
|
||||
r"(\d+(?:\.\d+)?)\s*(?:m?A(?:DC|RMS)?|amps?)\b", text, re.IGNORECASE
|
||||
)
|
||||
if currents:
|
||||
specs["currents_A"] = sorted(set(float(c) for c in currents))
|
||||
|
||||
# Temperature range
|
||||
temps = re.findall(
|
||||
r"(-?\d+)\s*(?:deg(?:ree)?s?\s*)?[°]?\s*C\b", text
|
||||
)
|
||||
if temps:
|
||||
t_vals = sorted(set(int(t) for t in temps))
|
||||
specs["temperature_range_C"] = {"min": t_vals[0], "max": t_vals[-1]}
|
||||
|
||||
# Package / footprint
|
||||
packages = re.findall(
|
||||
r"\b(SOT-?\d+|QFP-?\d+|QFN-?\d+|BGA-?\d+|DIP-?\d+|SOIC-?\d+|TSSOP-?\d+|TO-?\d+)\b",
|
||||
text, re.IGNORECASE,
|
||||
)
|
||||
if packages:
|
||||
specs["packages"] = sorted(set(p.upper() for p in packages))
|
||||
|
||||
# Part numbers (common patterns)
|
||||
parts = re.findall(
|
||||
r"\b([A-Z]{2,5}\d{3,}[A-Z0-9\-]*)\b", text
|
||||
)
|
||||
if parts:
|
||||
# Deduplicate and keep top 10
|
||||
seen = set()
|
||||
unique = []
|
||||
for p in parts:
|
||||
if p not in seen:
|
||||
seen.add(p)
|
||||
unique.append(p)
|
||||
specs["part_numbers"] = unique[:10]
|
||||
|
||||
# Frequency / clock
|
||||
freqs = re.findall(
|
||||
r"(\d+(?:\.\d+)?)\s*(MHz|GHz|kHz)\b", text, re.IGNORECASE
|
||||
)
|
||||
if freqs:
|
||||
specs["frequencies"] = [
|
||||
{"value": float(v), "unit": u} for v, u in freqs
|
||||
]
|
||||
|
||||
# Memory sizes
|
||||
mem = re.findall(
|
||||
r"(\d+(?:\.\d+)?)\s*(KB|MB|GB|kB)\b", text, re.IGNORECASE
|
||||
)
|
||||
if mem:
|
||||
specs["memory"] = [{"value": float(v), "unit": u.upper()} for v, u in mem]
|
||||
|
||||
return specs
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Process one PDF
|
||||
# ---------------------------------------------------------------------------
|
||||
def process_pdf(pdf_path: Path, backend: Optional[str] = None) -> Dict[str, Any]:
|
||||
"""Full pipeline: OCR -> text -> structured specs."""
|
||||
print(f" Processing {pdf_path.name} ...")
|
||||
text, used_backend = extract_text(pdf_path, backend)
|
||||
print(f" Backend: {used_backend} | Extracted {len(text)} chars")
|
||||
|
||||
specs = extract_specs(text, pdf_path.name)
|
||||
specs["_ocr_backend"] = used_backend
|
||||
specs["_text_length"] = len(text)
|
||||
# Optionally include raw text excerpt for debugging
|
||||
specs["_text_excerpt"] = text[:500] + ("..." if len(text) > 500 else "")
|
||||
return specs
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CLI
|
||||
# ---------------------------------------------------------------------------
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="OCR pipeline for component datasheets (free, local)",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog=__doc__,
|
||||
)
|
||||
parser.add_argument("--pdf", type=Path, help="Single PDF file")
|
||||
parser.add_argument("--dir", type=Path, help="Directory of PDFs (batch mode)")
|
||||
parser.add_argument("--output", type=Path, help="Output JSON file (single mode)")
|
||||
parser.add_argument("--output-dir", type=Path, help="Output directory (batch mode)")
|
||||
parser.add_argument(
|
||||
"--backend",
|
||||
choices=["marker", "marker_cli", "surya", "pypdf2"],
|
||||
help="Force a specific OCR backend (default: auto-fallback)",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.pdf:
|
||||
if not args.pdf.exists():
|
||||
sys.exit(f"ERROR: file not found: {args.pdf}")
|
||||
result = process_pdf(args.pdf, args.backend)
|
||||
out = args.output or Path(args.pdf.stem + "_specs.json")
|
||||
out.write_text(json.dumps(result, indent=2, ensure_ascii=False))
|
||||
print(f" -> {out}")
|
||||
|
||||
elif args.dir:
|
||||
if not args.dir.is_dir():
|
||||
sys.exit(f"ERROR: not a directory: {args.dir}")
|
||||
out_dir = args.output_dir or Path("specs")
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
pdfs = sorted(args.dir.glob("*.pdf"))
|
||||
if not pdfs:
|
||||
sys.exit(f"ERROR: no PDF files found in {args.dir}")
|
||||
|
||||
print(f" Found {len(pdfs)} PDFs in {args.dir}")
|
||||
all_specs: List[Dict[str, Any]] = []
|
||||
for pdf in pdfs:
|
||||
try:
|
||||
specs = process_pdf(pdf, args.backend)
|
||||
out_file = out_dir / (pdf.stem + "_specs.json")
|
||||
out_file.write_text(json.dumps(specs, indent=2, ensure_ascii=False))
|
||||
all_specs.append(specs)
|
||||
except Exception as exc:
|
||||
print(f" WARN: failed on {pdf.name}: {exc}")
|
||||
|
||||
# Summary file
|
||||
summary = out_dir / "_batch_summary.json"
|
||||
summary.write_text(json.dumps(all_specs, indent=2, ensure_ascii=False))
|
||||
print(f"\n Batch complete: {len(all_specs)}/{len(pdfs)} succeeded -> {out_dir}")
|
||||
|
||||
else:
|
||||
parser.print_help()
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,409 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
STT Pipeline — FREE alternative to paid speech-to-text APIs (T-MS-021).
|
||||
|
||||
Transcribes audio to text and extracts action items, using local engines
|
||||
with graceful fallback:
|
||||
|
||||
1. whisper.cpp (fastest, via subprocess — requires compiled binary)
|
||||
2. openai-whisper (Python, GPU-accelerated)
|
||||
3. vosk (lightweight, CPU-only, offline)
|
||||
|
||||
Usage:
|
||||
# Transcribe audio
|
||||
python3 stt_pipeline.py transcribe --audio meeting.wav --output transcript.md
|
||||
|
||||
# Extract action items from transcript
|
||||
python3 stt_pipeline.py actions --transcript transcript.md --output actions.json
|
||||
|
||||
# Full pipeline (transcribe + extract)
|
||||
python3 stt_pipeline.py full --audio meeting.wav --output-dir results/
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Transcription backends
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _find_whisper_cpp() -> Optional[str]:
|
||||
"""Locate whisper.cpp main binary."""
|
||||
import shutil
|
||||
|
||||
# Check common locations
|
||||
candidates = [
|
||||
os.environ.get("WHISPER_CPP_BIN", ""),
|
||||
"whisper-cpp",
|
||||
"main", # default binary name in whisper.cpp build
|
||||
str(Path.home() / "whisper.cpp" / "main"),
|
||||
"/usr/local/bin/whisper-cpp",
|
||||
]
|
||||
for c in candidates:
|
||||
if c and shutil.which(c):
|
||||
return c
|
||||
return None
|
||||
|
||||
|
||||
def _find_whisper_cpp_model() -> Optional[str]:
|
||||
"""Locate a whisper.cpp GGML model file."""
|
||||
search_dirs = [
|
||||
Path(os.environ.get("WHISPER_CPP_MODELS", "")),
|
||||
Path.home() / "whisper.cpp" / "models",
|
||||
Path.home() / ".cache" / "whisper-cpp",
|
||||
Path("/usr/local/share/whisper-cpp/models"),
|
||||
]
|
||||
preferred = ["ggml-base.en.bin", "ggml-base.bin", "ggml-small.en.bin", "ggml-small.bin"]
|
||||
|
||||
for d in search_dirs:
|
||||
if not d.is_dir():
|
||||
continue
|
||||
for name in preferred:
|
||||
p = d / name
|
||||
if p.exists():
|
||||
return str(p)
|
||||
# Any .bin file
|
||||
bins = list(d.glob("ggml-*.bin"))
|
||||
if bins:
|
||||
return str(bins[0])
|
||||
return None
|
||||
|
||||
|
||||
def transcribe_whisper_cpp(audio_path: Path, model_path: Optional[str] = None) -> str:
|
||||
"""Transcribe using whisper.cpp subprocess."""
|
||||
binary = _find_whisper_cpp()
|
||||
if not binary:
|
||||
raise RuntimeError(
|
||||
"whisper.cpp not found. Install from https://github.com/ggerganov/whisper.cpp\n"
|
||||
" or set WHISPER_CPP_BIN=/path/to/main"
|
||||
)
|
||||
|
||||
model = model_path or _find_whisper_cpp_model()
|
||||
if not model:
|
||||
raise RuntimeError(
|
||||
"No whisper.cpp model found. Download one:\n"
|
||||
" cd ~/whisper.cpp && bash ./models/download-ggml-model.sh base.en\n"
|
||||
" or set WHISPER_CPP_MODELS=/path/to/models/"
|
||||
)
|
||||
|
||||
# whisper.cpp expects 16kHz WAV; convert if needed
|
||||
wav_path = _ensure_wav_16k(audio_path)
|
||||
|
||||
with tempfile.NamedTemporaryFile(suffix=".txt", delete=False) as tmp:
|
||||
tmp_out = tmp.name
|
||||
|
||||
try:
|
||||
cmd = [binary, "-m", model, "-f", str(wav_path), "-otxt", "-of", tmp_out.replace(".txt", "")]
|
||||
result = subprocess.run(cmd, capture_output=True, text=True, timeout=600)
|
||||
if result.returncode != 0:
|
||||
raise RuntimeError(f"whisper.cpp failed: {result.stderr[:500]}")
|
||||
return Path(tmp_out).read_text(encoding="utf-8")
|
||||
finally:
|
||||
Path(tmp_out).unlink(missing_ok=True)
|
||||
if wav_path != audio_path:
|
||||
wav_path.unlink(missing_ok=True)
|
||||
|
||||
|
||||
def transcribe_openai_whisper(audio_path: Path, model_size: str = "base") -> str:
|
||||
"""Transcribe using openai-whisper Python package."""
|
||||
try:
|
||||
import whisper
|
||||
except ImportError:
|
||||
raise RuntimeError("openai-whisper not installed: pip install openai-whisper")
|
||||
|
||||
model = whisper.load_model(model_size)
|
||||
result = model.transcribe(str(audio_path))
|
||||
return result["text"]
|
||||
|
||||
|
||||
def transcribe_vosk(audio_path: Path, model_path: Optional[str] = None) -> str:
|
||||
"""Transcribe using Vosk (lightweight, CPU-only)."""
|
||||
try:
|
||||
from vosk import Model, KaldiRecognizer
|
||||
except ImportError:
|
||||
raise RuntimeError("vosk not installed: pip install vosk")
|
||||
|
||||
import wave
|
||||
|
||||
wav_path = _ensure_wav_16k(audio_path)
|
||||
|
||||
if model_path:
|
||||
model = Model(model_path)
|
||||
else:
|
||||
# Vosk auto-downloads a small model
|
||||
try:
|
||||
model = Model(lang="en-us")
|
||||
except Exception:
|
||||
raise RuntimeError(
|
||||
"No Vosk model found. Download from https://alphacephei.com/vosk/models\n"
|
||||
" or: pip install vosk (auto-downloads small model)"
|
||||
)
|
||||
|
||||
wf = wave.open(str(wav_path), "rb")
|
||||
if wf.getnchannels() != 1 or wf.getsampwidth() != 2 or wf.getframerate() != 16000:
|
||||
wf.close()
|
||||
raise RuntimeError("Audio must be mono 16kHz 16-bit WAV for Vosk")
|
||||
|
||||
rec = KaldiRecognizer(model, 16000)
|
||||
rec.SetWords(True)
|
||||
|
||||
texts: List[str] = []
|
||||
while True:
|
||||
data = wf.readframes(4000)
|
||||
if len(data) == 0:
|
||||
break
|
||||
if rec.AcceptWaveform(data):
|
||||
res = json.loads(rec.Result())
|
||||
if res.get("text"):
|
||||
texts.append(res["text"])
|
||||
|
||||
final = json.loads(rec.FinalResult())
|
||||
if final.get("text"):
|
||||
texts.append(final["text"])
|
||||
|
||||
wf.close()
|
||||
if wav_path != audio_path:
|
||||
wav_path.unlink(missing_ok=True)
|
||||
|
||||
return " ".join(texts)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Audio format helper
|
||||
# ---------------------------------------------------------------------------
|
||||
def _ensure_wav_16k(audio_path: Path) -> Path:
|
||||
"""Convert audio to 16kHz mono WAV if needed (requires ffmpeg)."""
|
||||
import shutil
|
||||
|
||||
if audio_path.suffix.lower() == ".wav":
|
||||
# Quick check — might already be 16k mono
|
||||
return audio_path
|
||||
|
||||
if not shutil.which("ffmpeg"):
|
||||
print(" WARN: ffmpeg not found, hoping input is already valid WAV")
|
||||
return audio_path
|
||||
|
||||
tmp = Path(tempfile.mktemp(suffix=".wav"))
|
||||
subprocess.run(
|
||||
["ffmpeg", "-y", "-i", str(audio_path), "-ar", "16000", "-ac", "1", "-f", "wav", str(tmp)],
|
||||
capture_output=True,
|
||||
timeout=120,
|
||||
)
|
||||
return tmp
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Transcription dispatcher with fallback
|
||||
# ---------------------------------------------------------------------------
|
||||
BACKENDS = [
|
||||
("whisper_cpp", transcribe_whisper_cpp),
|
||||
("openai_whisper", transcribe_openai_whisper),
|
||||
("vosk", transcribe_vosk),
|
||||
]
|
||||
|
||||
|
||||
def transcribe(audio_path: Path, backend: Optional[str] = None) -> Tuple[str, str]:
|
||||
"""Transcribe audio. Returns (text, backend_used)."""
|
||||
if backend:
|
||||
for name, fn in BACKENDS:
|
||||
if name == backend:
|
||||
return fn(audio_path), name
|
||||
sys.exit(f"ERROR: unknown backend '{backend}'. Choose from: {[b[0] for b in BACKENDS]}")
|
||||
|
||||
errors: List[str] = []
|
||||
for name, fn in BACKENDS:
|
||||
try:
|
||||
text = fn(audio_path)
|
||||
if text.strip():
|
||||
return text, name
|
||||
errors.append(f"{name}: empty output")
|
||||
except Exception as exc:
|
||||
errors.append(f"{name}: {exc}")
|
||||
|
||||
sys.exit(
|
||||
f"ERROR: all STT backends failed for {audio_path.name}:\n"
|
||||
+ "\n".join(f" - {e}" for e in errors)
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Action item extraction (regex-based, no LLM)
|
||||
# ---------------------------------------------------------------------------
|
||||
ACTION_PATTERNS = [
|
||||
# "TODO: ..."
|
||||
re.compile(r"(?:TODO|FIXME|ACTION|A[Cc]tion\s*[Ii]tem)[:\s]+(.+?)(?:\.|$)", re.MULTILINE),
|
||||
# "we need to ..."
|
||||
re.compile(r"(?:we\s+(?:need|should|must|have)\s+to|il\s+faut)\s+(.+?)(?:\.|$)", re.IGNORECASE),
|
||||
# "X will ..." / "X va ..."
|
||||
re.compile(r"(\w+)\s+(?:will|va|doit)\s+(.+?)(?:\.|$)", re.IGNORECASE),
|
||||
# "let's ..." / "on va ..."
|
||||
re.compile(r"(?:let'?s|on\s+va)\s+(.+?)(?:\.|$)", re.IGNORECASE),
|
||||
# Deadline patterns
|
||||
re.compile(r"(?:deadline|before|by|avant)\s+(\w+\s+\d+|\d{4}-\d{2}-\d{2})", re.IGNORECASE),
|
||||
]
|
||||
|
||||
|
||||
def extract_actions(text: str) -> List[Dict[str, str]]:
|
||||
"""Extract action items from transcript text."""
|
||||
actions: List[Dict[str, str]] = []
|
||||
seen: set = set()
|
||||
|
||||
for pattern in ACTION_PATTERNS:
|
||||
for match in pattern.finditer(text):
|
||||
full = match.group(0).strip()
|
||||
if full and full not in seen:
|
||||
seen.add(full)
|
||||
actions.append({
|
||||
"action": full,
|
||||
"context": _get_context(text, match.start(), window=100),
|
||||
})
|
||||
|
||||
return actions
|
||||
|
||||
|
||||
def _get_context(text: str, pos: int, window: int = 100) -> str:
|
||||
"""Get surrounding text for context."""
|
||||
start = max(0, pos - window)
|
||||
end = min(len(text), pos + window)
|
||||
return text[start:end].strip()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Output formatters
|
||||
# ---------------------------------------------------------------------------
|
||||
def format_transcript_md(text: str, audio_name: str, backend: str) -> str:
|
||||
"""Format transcript as markdown."""
|
||||
now = datetime.now().strftime("%Y-%m-%d %H:%M")
|
||||
return (
|
||||
f"# Transcript: {audio_name}\n\n"
|
||||
f"- **Date**: {now}\n"
|
||||
f"- **Backend**: {backend}\n"
|
||||
f"- **Length**: {len(text)} chars\n\n"
|
||||
f"---\n\n"
|
||||
f"{text}\n"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CLI commands
|
||||
# ---------------------------------------------------------------------------
|
||||
def cmd_transcribe(args):
|
||||
"""Transcribe audio to text."""
|
||||
if not args.audio.exists():
|
||||
sys.exit(f"ERROR: file not found: {args.audio}")
|
||||
|
||||
print(f" Transcribing {args.audio.name} ...")
|
||||
text, backend = transcribe(args.audio, getattr(args, "backend", None))
|
||||
print(f" Backend: {backend} | {len(text)} chars")
|
||||
|
||||
md = format_transcript_md(text, args.audio.name, backend)
|
||||
out = args.output or Path(args.audio.stem + "_transcript.md")
|
||||
out.write_text(md, encoding="utf-8")
|
||||
print(f" -> {out}")
|
||||
|
||||
|
||||
def cmd_actions(args):
|
||||
"""Extract action items from a transcript."""
|
||||
if not args.transcript.exists():
|
||||
sys.exit(f"ERROR: file not found: {args.transcript}")
|
||||
|
||||
text = args.transcript.read_text(encoding="utf-8")
|
||||
print(f" Extracting actions from {args.transcript.name} ({len(text)} chars) ...")
|
||||
|
||||
actions = extract_actions(text)
|
||||
print(f" Found {len(actions)} action items")
|
||||
|
||||
result = {
|
||||
"source": str(args.transcript),
|
||||
"extracted_at": datetime.now().isoformat(),
|
||||
"actions": actions,
|
||||
}
|
||||
|
||||
out = args.output or Path(args.transcript.stem + "_actions.json")
|
||||
out.write_text(json.dumps(result, indent=2, ensure_ascii=False))
|
||||
print(f" -> {out}")
|
||||
|
||||
|
||||
def cmd_full(args):
|
||||
"""Full pipeline: transcribe + extract actions."""
|
||||
if not args.audio.exists():
|
||||
sys.exit(f"ERROR: file not found: {args.audio}")
|
||||
|
||||
out_dir = args.output_dir or Path(".")
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Transcribe
|
||||
print(f" Transcribing {args.audio.name} ...")
|
||||
text, backend = transcribe(args.audio, getattr(args, "backend", None))
|
||||
print(f" Backend: {backend} | {len(text)} chars")
|
||||
|
||||
md = format_transcript_md(text, args.audio.name, backend)
|
||||
transcript_path = out_dir / (args.audio.stem + "_transcript.md")
|
||||
transcript_path.write_text(md, encoding="utf-8")
|
||||
print(f" -> {transcript_path}")
|
||||
|
||||
# Actions
|
||||
actions = extract_actions(text)
|
||||
print(f" Found {len(actions)} action items")
|
||||
|
||||
result = {
|
||||
"source": str(args.audio),
|
||||
"backend": backend,
|
||||
"extracted_at": datetime.now().isoformat(),
|
||||
"actions": actions,
|
||||
}
|
||||
actions_path = out_dir / (args.audio.stem + "_actions.json")
|
||||
actions_path.write_text(json.dumps(result, indent=2, ensure_ascii=False))
|
||||
print(f" -> {actions_path}")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CLI
|
||||
# ---------------------------------------------------------------------------
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="STT pipeline — free, local speech-to-text + action extraction",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog=__doc__,
|
||||
)
|
||||
sub = parser.add_subparsers(dest="command")
|
||||
|
||||
# transcribe
|
||||
p_tr = sub.add_parser("transcribe", help="Transcribe audio to text")
|
||||
p_tr.add_argument("--audio", type=Path, required=True)
|
||||
p_tr.add_argument("--output", type=Path)
|
||||
p_tr.add_argument("--backend", choices=["whisper_cpp", "openai_whisper", "vosk"])
|
||||
p_tr.set_defaults(func=cmd_transcribe)
|
||||
|
||||
# actions
|
||||
p_ac = sub.add_parser("actions", help="Extract action items from transcript")
|
||||
p_ac.add_argument("--transcript", type=Path, required=True)
|
||||
p_ac.add_argument("--output", type=Path)
|
||||
p_ac.set_defaults(func=cmd_actions)
|
||||
|
||||
# full
|
||||
p_fu = sub.add_parser("full", help="Transcribe + extract actions")
|
||||
p_fu.add_argument("--audio", type=Path, required=True)
|
||||
p_fu.add_argument("--output-dir", type=Path)
|
||||
p_fu.add_argument("--backend", choices=["whisper_cpp", "openai_whisper", "vosk"])
|
||||
p_fu.set_defaults(func=cmd_full)
|
||||
|
||||
args = parser.parse_args()
|
||||
if not args.command:
|
||||
parser.print_help()
|
||||
sys.exit(1)
|
||||
|
||||
args.func(args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,18 @@
|
||||
# Ollama Modelfile — auto-generated by local_finetune.py
|
||||
# Base: {{BASE_MODEL}}
|
||||
# Name: {{MODEL_NAME}}
|
||||
|
||||
FROM {{GGUF_PATH}}
|
||||
|
||||
PARAMETER temperature 0.2
|
||||
PARAMETER top_p 0.9
|
||||
PARAMETER repeat_penalty 1.1
|
||||
|
||||
TEMPLATE """<|im_start|>system
|
||||
{{ .System }}<|im_end|>
|
||||
<|im_start|>user
|
||||
{{ .Prompt }}<|im_end|>
|
||||
<|im_start|>assistant
|
||||
"""
|
||||
|
||||
SYSTEM """You are a specialised engineering assistant fine-tuned on KXKM hardware project data (KiCad, embedded, DSP, power electronics). Answer precisely and concisely. When generating KiCad artifacts, use valid S-expression syntax."""
|
||||
@@ -0,0 +1,371 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Local fine-tuning via Unsloth + QLoRA — FREE alternative to Mistral paid fine-tune API.
|
||||
|
||||
Target hardware: KXKM RTX 4090 (24 GB VRAM).
|
||||
Supported bases: Mistral-7B-v0.3, Codestral-22B-v0.1 (HuggingFace weights).
|
||||
|
||||
Usage:
|
||||
python3 local_finetune.py \
|
||||
--base mistral-7b \
|
||||
--dataset datasets/kicad_merged.jsonl \
|
||||
--output models/kicad_qlora
|
||||
|
||||
python3 local_finetune.py \
|
||||
--base codestral-22b \
|
||||
--dataset datasets/embedded/stm32_merged.jsonl \
|
||||
--output models/embedded_qlora \
|
||||
--steps 200 --lr 3e-5
|
||||
|
||||
# Export to GGUF for Ollama after training:
|
||||
python3 local_finetune.py \
|
||||
--export-gguf models/kicad_qlora \
|
||||
--ollama-name mascarade-kicad
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Model registry — maps friendly names to HuggingFace repo IDs
|
||||
# ---------------------------------------------------------------------------
|
||||
BASE_MODELS: Dict[str, str] = {
|
||||
"mistral-7b": "mistralai/Mistral-7B-v0.3",
|
||||
"codestral-22b": "mistralai/Codestral-22B-v0.1",
|
||||
}
|
||||
|
||||
# QLoRA defaults (4-bit, rank 16)
|
||||
QLORA_DEFAULTS = dict(
|
||||
r=16,
|
||||
lora_alpha=16,
|
||||
lora_dropout=0.0,
|
||||
target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
|
||||
bias="none",
|
||||
task_type="CAUSAL_LM",
|
||||
)
|
||||
|
||||
TEMPLATE_DIR = Path(__file__).resolve().parent
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Dataset helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
def load_jsonl(path: Path) -> List[Dict[str, Any]]:
|
||||
"""Load a JSONL file (one JSON object per line)."""
|
||||
records: List[Dict[str, Any]] = []
|
||||
with open(path, "r", encoding="utf-8") as fh:
|
||||
for lineno, line in enumerate(fh, 1):
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
records.append(json.loads(line))
|
||||
except json.JSONDecodeError as exc:
|
||||
print(f" WARN: skipping line {lineno} in {path.name}: {exc}")
|
||||
return records
|
||||
|
||||
|
||||
def prepare_dataset(path: Path):
|
||||
"""Return a HuggingFace Dataset from a JSONL file.
|
||||
|
||||
Expected JSONL format (chat-style):
|
||||
{"messages": [{"role":"user","content":"..."}, {"role":"assistant","content":"..."}]}
|
||||
OR simple instruction/output:
|
||||
{"instruction": "...", "output": "..."}
|
||||
"""
|
||||
try:
|
||||
from datasets import Dataset
|
||||
except ImportError:
|
||||
sys.exit("ERROR: pip install datasets (required for training)")
|
||||
|
||||
records = load_jsonl(path)
|
||||
if not records:
|
||||
sys.exit(f"ERROR: dataset {path} is empty")
|
||||
|
||||
# Normalise to a single 'text' column using ChatML formatting
|
||||
texts: List[str] = []
|
||||
for rec in records:
|
||||
if "messages" in rec:
|
||||
parts = []
|
||||
for msg in rec["messages"]:
|
||||
role = msg.get("role", "user")
|
||||
content = msg.get("content", "")
|
||||
parts.append(f"<|im_start|>{role}\n{content}<|im_end|>")
|
||||
texts.append("\n".join(parts))
|
||||
elif "instruction" in rec:
|
||||
texts.append(
|
||||
f"<|im_start|>user\n{rec['instruction']}<|im_end|>\n"
|
||||
f"<|im_start|>assistant\n{rec.get('output', '')}<|im_end|>"
|
||||
)
|
||||
elif "text" in rec:
|
||||
texts.append(rec["text"])
|
||||
else:
|
||||
# Best-effort: serialise the whole object
|
||||
texts.append(json.dumps(rec, ensure_ascii=False))
|
||||
|
||||
print(f" Loaded {len(texts)} training examples from {path.name}")
|
||||
return Dataset.from_dict({"text": texts})
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Training
|
||||
# ---------------------------------------------------------------------------
|
||||
def train(
|
||||
base_name: str,
|
||||
dataset_path: Path,
|
||||
output_dir: Path,
|
||||
max_steps: int = 100,
|
||||
lr: float = 2e-5,
|
||||
batch_size: int = 4,
|
||||
grad_accum: int = 4,
|
||||
max_seq_length: int = 2048,
|
||||
):
|
||||
"""Run QLoRA fine-tuning with Unsloth (fast LoRA for consumer GPUs)."""
|
||||
|
||||
# --- Lazy imports so the script can still show --help without GPU libs ---
|
||||
try:
|
||||
from unsloth import FastLanguageModel
|
||||
except ImportError:
|
||||
sys.exit(
|
||||
"ERROR: Unsloth not installed.\n"
|
||||
" pip install 'unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git'\n"
|
||||
" pip install --no-deps trl peft accelerate bitsandbytes"
|
||||
)
|
||||
try:
|
||||
from trl import SFTTrainer
|
||||
from transformers import TrainingArguments
|
||||
except ImportError:
|
||||
sys.exit("ERROR: pip install trl transformers")
|
||||
|
||||
model_id = BASE_MODELS.get(base_name)
|
||||
if model_id is None:
|
||||
sys.exit(
|
||||
f"ERROR: unknown base '{base_name}'. "
|
||||
f"Choose from: {', '.join(BASE_MODELS)}"
|
||||
)
|
||||
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# 1. Load base model in 4-bit
|
||||
print(f"\n==> Loading {model_id} in 4-bit quantisation ...")
|
||||
model, tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name=model_id,
|
||||
max_seq_length=max_seq_length,
|
||||
dtype=None, # auto
|
||||
load_in_4bit=True,
|
||||
)
|
||||
|
||||
# 2. Attach LoRA adapters
|
||||
print("==> Attaching QLoRA adapters ...")
|
||||
model = FastLanguageModel.get_peft_model(
|
||||
model,
|
||||
r=QLORA_DEFAULTS["r"],
|
||||
lora_alpha=QLORA_DEFAULTS["lora_alpha"],
|
||||
lora_dropout=QLORA_DEFAULTS["lora_dropout"],
|
||||
target_modules=QLORA_DEFAULTS["target_modules"],
|
||||
bias=QLORA_DEFAULTS["bias"],
|
||||
)
|
||||
|
||||
# 3. Prepare dataset
|
||||
print(f"==> Preparing dataset from {dataset_path} ...")
|
||||
ds = prepare_dataset(dataset_path)
|
||||
|
||||
# 4. Train
|
||||
print(f"==> Training for {max_steps} steps (lr={lr}, bs={batch_size}x{grad_accum}) ...")
|
||||
training_args = TrainingArguments(
|
||||
output_dir=str(output_dir / "checkpoints"),
|
||||
max_steps=max_steps,
|
||||
learning_rate=lr,
|
||||
per_device_train_batch_size=batch_size,
|
||||
gradient_accumulation_steps=grad_accum,
|
||||
fp16=True,
|
||||
logging_steps=10,
|
||||
save_steps=max_steps, # save at the end
|
||||
warmup_steps=min(10, max_steps // 10),
|
||||
optim="adamw_8bit",
|
||||
seed=42,
|
||||
report_to="none",
|
||||
)
|
||||
|
||||
trainer = SFTTrainer(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
train_dataset=ds,
|
||||
dataset_text_field="text",
|
||||
max_seq_length=max_seq_length,
|
||||
args=training_args,
|
||||
)
|
||||
|
||||
trainer.train()
|
||||
|
||||
# 5. Save adapter weights
|
||||
adapter_dir = output_dir / "adapter"
|
||||
print(f"==> Saving LoRA adapter to {adapter_dir} ...")
|
||||
model.save_pretrained(str(adapter_dir))
|
||||
tokenizer.save_pretrained(str(adapter_dir))
|
||||
|
||||
# 6. Save training metadata
|
||||
meta = {
|
||||
"base_model": model_id,
|
||||
"base_alias": base_name,
|
||||
"dataset": str(dataset_path),
|
||||
"qlora": QLORA_DEFAULTS,
|
||||
"training": {
|
||||
"max_steps": max_steps,
|
||||
"lr": lr,
|
||||
"batch_size": batch_size,
|
||||
"grad_accum": grad_accum,
|
||||
"max_seq_length": max_seq_length,
|
||||
},
|
||||
}
|
||||
meta_path = output_dir / "training_meta.json"
|
||||
meta_path.write_text(json.dumps(meta, indent=2, ensure_ascii=False))
|
||||
print(f"==> Metadata saved to {meta_path}")
|
||||
|
||||
print("\n==> Training complete!")
|
||||
print(f" Adapter: {adapter_dir}")
|
||||
print(f" Metadata: {meta_path}")
|
||||
print(
|
||||
f"\n Next step — export to GGUF for Ollama:\n"
|
||||
f" python3 {Path(__file__).name} --export-gguf {output_dir} --ollama-name mascarade-kicad"
|
||||
)
|
||||
return output_dir
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# GGUF export + Ollama import
|
||||
# ---------------------------------------------------------------------------
|
||||
def export_gguf(model_dir: Path, ollama_name: Optional[str] = None):
|
||||
"""Merge LoRA adapter back into base, quantise to GGUF, optionally register in Ollama."""
|
||||
try:
|
||||
from unsloth import FastLanguageModel
|
||||
except ImportError:
|
||||
sys.exit("ERROR: Unsloth not installed (needed for GGUF export).")
|
||||
|
||||
meta_path = model_dir / "training_meta.json"
|
||||
if not meta_path.exists():
|
||||
sys.exit(f"ERROR: {meta_path} not found — is this a local_finetune output dir?")
|
||||
meta = json.loads(meta_path.read_text())
|
||||
|
||||
adapter_dir = model_dir / "adapter"
|
||||
gguf_dir = model_dir / "gguf"
|
||||
gguf_dir.mkdir(exist_ok=True)
|
||||
|
||||
print(f"==> Loading base model + adapter from {adapter_dir} ...")
|
||||
model, tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name=str(adapter_dir),
|
||||
max_seq_length=meta["training"]["max_seq_length"],
|
||||
dtype=None,
|
||||
load_in_4bit=True,
|
||||
)
|
||||
|
||||
# Unsloth provides a convenient save_pretrained_gguf helper
|
||||
print(f"==> Exporting merged GGUF to {gguf_dir} ...")
|
||||
model.save_pretrained_gguf(
|
||||
str(gguf_dir),
|
||||
tokenizer,
|
||||
quantization_method="q4_k_m", # good quality / size trade-off
|
||||
)
|
||||
|
||||
gguf_files = list(gguf_dir.glob("*.gguf"))
|
||||
if not gguf_files:
|
||||
print("WARN: No .gguf file produced — check Unsloth version.")
|
||||
return
|
||||
|
||||
gguf_path = gguf_files[0]
|
||||
print(f"==> GGUF ready: {gguf_path}")
|
||||
|
||||
# Generate Modelfile from template
|
||||
template_path = TEMPLATE_DIR / "Modelfile.template"
|
||||
modelfile_path = model_dir / "Modelfile"
|
||||
if template_path.exists():
|
||||
content = template_path.read_text()
|
||||
content = content.replace("{{GGUF_PATH}}", str(gguf_path))
|
||||
content = content.replace("{{MODEL_NAME}}", ollama_name or model_dir.name)
|
||||
content = content.replace("{{BASE_MODEL}}", meta.get("base_model", "unknown"))
|
||||
modelfile_path.write_text(content)
|
||||
print(f"==> Modelfile written: {modelfile_path}")
|
||||
else:
|
||||
# Inline fallback
|
||||
modelfile_path.write_text(
|
||||
f"FROM {gguf_path}\n"
|
||||
f"PARAMETER temperature 0.2\n"
|
||||
f"PARAMETER top_p 0.9\n"
|
||||
f"SYSTEM You are a specialised engineering assistant fine-tuned on KXKM hardware data.\n"
|
||||
)
|
||||
print(f"==> Modelfile written (inline): {modelfile_path}")
|
||||
|
||||
# Optionally register in Ollama
|
||||
if ollama_name and shutil.which("ollama"):
|
||||
print(f"==> Registering in Ollama as '{ollama_name}' ...")
|
||||
result = subprocess.run(
|
||||
["ollama", "create", ollama_name, "-f", str(modelfile_path)],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
if result.returncode == 0:
|
||||
print(f" OK — run: ollama run {ollama_name}")
|
||||
else:
|
||||
print(f" WARN: ollama create failed: {result.stderr.strip()}")
|
||||
elif ollama_name:
|
||||
print(f" Ollama CLI not found. Import manually:\n ollama create {ollama_name} -f {modelfile_path}")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CLI
|
||||
# ---------------------------------------------------------------------------
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Local QLoRA fine-tuning (Unsloth) — free Mistral alternative",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog=__doc__,
|
||||
)
|
||||
|
||||
# Training mode
|
||||
parser.add_argument("--base", choices=list(BASE_MODELS), default="mistral-7b",
|
||||
help="Base model alias (default: mistral-7b)")
|
||||
parser.add_argument("--dataset", type=Path,
|
||||
help="Path to JSONL training dataset")
|
||||
parser.add_argument("--output", type=Path, default=Path("models/finetune_qlora"),
|
||||
help="Output directory for adapter + GGUF")
|
||||
parser.add_argument("--steps", type=int, default=100, help="Max training steps")
|
||||
parser.add_argument("--lr", type=float, default=2e-5, help="Learning rate")
|
||||
parser.add_argument("--batch-size", type=int, default=4)
|
||||
parser.add_argument("--grad-accum", type=int, default=4)
|
||||
parser.add_argument("--max-seq-length", type=int, default=2048)
|
||||
|
||||
# Export mode
|
||||
parser.add_argument("--export-gguf", type=Path, metavar="MODEL_DIR",
|
||||
help="Export an existing adapter dir to GGUF")
|
||||
parser.add_argument("--ollama-name", type=str,
|
||||
help="Register the GGUF model in Ollama with this name")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.export_gguf:
|
||||
export_gguf(args.export_gguf, args.ollama_name)
|
||||
elif args.dataset:
|
||||
train(
|
||||
base_name=args.base,
|
||||
dataset_path=args.dataset,
|
||||
output_dir=args.output,
|
||||
max_steps=args.steps,
|
||||
lr=args.lr,
|
||||
batch_size=args.batch_size,
|
||||
grad_accum=args.grad_accum,
|
||||
max_seq_length=args.max_seq_length,
|
||||
)
|
||||
else:
|
||||
parser.print_help()
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user