File size: 4,273 Bytes
9d614f1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 | """PUBLIC INTERFACE of Serving_Module. Other modules may import ONLY this file.
resolve_trained_dir(cfg, fetch_remote) -> Path local trained_model/ first, else OUTPUT_HF_REPO via `fetch_remote`
prepare_model(model, trained_dir, meta) -> model attaches the LoRA adapter when the checkpoint is adapter-based
build_engine(model, tokenizer, cfg) -> engine
start(engine, cfg) -> None Telegram bot if TELEGRAM_BOT_TOKEN is set, else CLI chat
Weights are loaded by the orchestrator (hardware tier logic lives in Hardware_Module); `fetch_remote` is injected
for the same reason, so this module never imports another module.
"""
from __future__ import annotations
from pathlib import Path
from typing import Callable, Dict, Optional
from .internal.cli_chat_window import run_cli
from .internal.inference_engine import InferenceEngine
from .internal.telegram_gateway import TelegramGateway
PLACEHOLDER_MARK = "PUT_YOUR_TOKEN"
def _has_meta(folder: Path) -> bool:
return (folder / "rss_meta.json").exists()
def resolve_trained_dir(cfg: Dict, fetch_remote: Optional[Callable] = None) -> Path:
root = Path(cfg["root"])
local = root / "trained_model"
if _has_meta(local):
print(f"[Serving] using local trained model: {local}")
return local
repo, token = cfg.get("output_hf_repo"), cfg.get("hf_token")
if repo and fetch_remote is not None:
try:
print(f"[Serving] no local trained model - pulling {repo} from Hugging Face ...")
token = None if (not token or PLACEHOLDER_MARK in token) else token
return Path(fetch_remote(repo, root / "trained_model", token))
except Exception as exc:
raise FileNotFoundError(
f"no trained model found locally ({local}) and could not download {repo}: "
f"{type(exc).__name__}: {exc}. Run: python RSS_AI_Engine/main.py --mode train_all") from exc
raise FileNotFoundError(f"no trained model found in {local}. Run: python RSS_AI_Engine/main.py --mode train_all")
def prepare_model(model, trained_dir, meta: Dict):
"""Attach the LoRA adapter for adapter-based checkpoints; fold it into the weights when that is possible."""
if not meta.get("has_adapter"):
return model
from peft import PeftModel
adapter = Path(trained_dir) / (meta.get("adapter_dir") or "adapter")
print(f"[Serving] attaching LoRA adapter {adapter}")
wrapped = PeftModel.from_pretrained(model, str(adapter), is_trainable=False)
plan = getattr(model, "_rss_plan", {}) or {}
ring = getattr(model, "_rss_ring", None)
if ring is None and not plan.get("quant"):
try:
merged = wrapped.merge_and_unload() # faster generation: no adapter matmuls
for attr in ("_rss_plan", "_rss_ring"):
setattr(merged, attr, getattr(model, attr, None))
return merged
except Exception as exc:
print(f"[Serving] adapter merge skipped ({type(exc).__name__}); serving with the adapter attached")
return wrapped
def build_engine(model, tokenizer, cfg: Dict) -> InferenceEngine:
return InferenceEngine(model, tokenizer, chat_file=cfg.get("chat_file"),
max_new_tokens=(cfg.get("tune") or {}).get("serve_max_new_tokens"),
active_vocab_file=cfg.get("active_vocab_file"))
def start(engine: InferenceEngine, cfg: Dict, telegram_session=None) -> str:
"""Blocks until the user quits. Returns which front-end ran: 'telegram' or 'cli'."""
token = (cfg.get("telegram_token") or "").strip()
if token and PLACEHOLDER_MARK not in token:
gateway = TelegramGateway(token, cfg.get("telegram_chat_id", ""), engine, session=telegram_session)
if gateway.verify():
gateway.notify_start()
try:
gateway.run()
except KeyboardInterrupt:
print("\n[Serving] Telegram bot stopped.")
return "telegram"
print("[Serving] Telegram token rejected -> falling back to the CLI chat window")
else:
print("[Serving] TELEGRAM_BOT_TOKEN is empty -> opening the interactive CLI chat window")
run_cli(engine)
return "cli"
|