File size: 4,297 Bytes
5ef3875 | 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 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 | """Local provider: YCode's own from-scratch model (``ycode-lm``), no API key.
A small model trained on a laptop can answer coding questions and write
short functions, but it has not learned the tool-calling protocol, so it
runs in answer-only mode: it never reads, edits or runs anything.
"""
from __future__ import annotations
from pathlib import Path
from ycode.errors import ConfigError
from ycode.llm.base import (
AssistantMessage,
LLMProvider,
LLMResponse,
Message,
StopReason,
StreamCallbacks,
ToolSpec,
Usage,
UserMessage,
)
_NOTE_PREFIX = "[Note:"
class LocalProvider(LLMProvider):
name = "local"
supports_tools = False
def __init__(self, model_dir: Path, *, max_new_tokens: int = 400, temperature: float = 0.7,
device: str = "auto", lm=None) -> None:
super().__init__(str(model_dir))
self.model_dir = Path(model_dir)
self.max_new_tokens = max_new_tokens
self.temperature = temperature
self.lm = lm if lm is not None else self._load(device)
from ycode.lm.assist import Assistant
# Big models are slow on a CPU, so they get fewer tries at writing a checked function.
big = getattr(self.lm, "num_params", 0) > 500e6
self.assistant = Assistant(self.lm, max_new_tokens=max_new_tokens, temperature=temperature,
candidates=2 if big else 4)
def _load(self, device: str):
try:
from ycode.lm.generate import LocalLM
except ImportError:
raise ConfigError(
"The local model needs PyTorch.",
hint='Install it with: pip install -e ".[local]"',
) from None
try:
return LocalLM(self.model_dir, device=device)
except FileNotFoundError:
from ycode.config import bundled_models
available = ", ".join(m.removeprefix("ycode-lm-") for m in bundled_models())
raise ConfigError(
f"No model found at {self.model_dir}.",
hint=(f"Bundled models available: {available} (e.g. `ycode --local v4`)." if available else
"YCode's trained models live in the repository's models/ folder: install from a clone "
"(git clone https://github.com/yazhiniyanMS/Code-agent && pip install -e .), "
"or set YCODE_MODELS_DIR to a folder containing them."),
) from None
except (ValueError, RuntimeError) as exc:
raise ConfigError(f"Could not load the local model at {self.model_dir}: {exc}") from None
@property
def display_name(self) -> str:
params = getattr(self.lm, "num_params", 0)
version = getattr(self.lm, "version", None)
size = f"{params / 1e6:.1f}M params, " if params else ""
name = f"YCode-LM v{version}" if isinstance(version, int) else "YCode-LM"
return f"{name} ({size}local, {self.model_dir})"
@staticmethod
def _turns(messages: list[Message]) -> list[tuple[str, str]]:
turns: list[tuple[str, str]] = []
for msg in messages:
if isinstance(msg, UserMessage):
text = msg.text
if text.startswith(_NOTE_PREFIX) and "\n\n" in text:
text = text.split("\n\n", 1)[1]
turns.append(("user", text))
elif isinstance(msg, AssistantMessage) and msg.text:
turns.append(("assistant", msg.text))
return turns
def complete(self, *, system: str, messages: list[Message], tools: list[ToolSpec],
callbacks: StreamCallbacks | None = None) -> LLMResponse:
callbacks = callbacks or StreamCallbacks()
turns = self._turns(messages)[-6:]
if not turns or turns[-1][0] != "user":
turns.append(("user", "Continue."))
answer = self.assistant.answer(turns, on_text=callbacks.text)
if not answer.strip():
answer = "(The local model did not produce an answer. Try rephrasing the question.)"
callbacks.text(answer)
return LLMResponse(
AssistantMessage(text=answer, provider=self.name),
StopReason.END_TURN,
Usage(output_tokens=len(answer) // 4),
)
|