| |
| """Generate once using local Ollama or a local Unsloth adapter. No paid API.""" |
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import importlib.metadata |
| import json |
| from pathlib import Path |
| import sys |
| import time |
| import urllib.parse |
| import urllib.request |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src")) |
| from vibegame.project import SYSTEM_PROMPT, parse_project |
|
|
|
|
| def file_hash(path: Path) -> str: |
| digest = hashlib.sha256() |
| with path.open("rb") as handle: |
| for block in iter(lambda: handle.read(8 * 1024 * 1024), b""): |
| digest.update(block) |
| return digest.hexdigest() |
|
|
|
|
| def directory_fingerprint(path: Path) -> dict: |
| files = {str(item.relative_to(path)): file_hash(item) for item in sorted(path.rglob("*")) |
| if item.is_file() and item.suffix in {".json", ".safetensors", ".bin", ".model", ".gguf"}} |
| return {"path": str(path.resolve()), "sha256": hashlib.sha256(json.dumps(files, sort_keys=True).encode()).hexdigest(), "files": files} |
|
|
|
|
| class OllamaGenerator: |
| def __init__(self, model: str, url: str = "http://127.0.0.1:11434", **kwargs): |
| parsed = urllib.parse.urlparse(url) |
| if parsed.scheme != "http" or parsed.hostname not in {"127.0.0.1", "localhost", "::1"} or parsed.username: |
| raise ValueError("Only local, unauthenticated Ollama HTTP is allowed") |
| self.url, self.model = url.rstrip("/"), model |
| tags = self._post("/api/tags", None) |
| matching = [item for item in tags.get("models", []) if item.get("name") in {model, model + ":latest"} or item.get("model") in {model, model + ":latest"}] |
| if not matching: |
| raise RuntimeError("Model is not installed in local Ollama; automatic pull is disabled") |
| self.metadata = {"backend": "ollama-local", "model": model, "model_digest": matching[0]["digest"], |
| "runtime": self._post("/api/version", None), "system_prompt_sha256": hashlib.sha256(SYSTEM_PROMPT.encode()).hexdigest()} |
|
|
| def _post(self, endpoint: str, payload): |
| body = None if payload is None else json.dumps(payload).encode() |
| request = urllib.request.Request(self.url + endpoint, data=body, headers={"Content-Type": "application/json"}) |
| with urllib.request.urlopen(request, timeout=900) as response: |
| return json.load(response) |
|
|
| def generate(self, prompt: str, max_new_tokens: int = 3072) -> dict: |
| start = time.monotonic() |
| result = self._post("/api/generate", {"model": self.model, "system": SYSTEM_PROMPT, "prompt": prompt, |
| "stream": False, "options": {"temperature": 0, "seed": 42, "num_predict": max_new_tokens, "num_ctx": 4096}}) |
| elapsed = time.monotonic() - start |
| return {"output": result["response"], "generation_seconds": elapsed, "generation_call_seconds": elapsed, |
| "tokens": result.get("eval_count"), "done_reason": result.get("done_reason")} |
|
|
|
|
| class UnslothGenerator: |
| def __init__(self, adapter: str, base_model: str | None = None, **kwargs): |
| from unsloth import FastLanguageModel |
| import torch |
| if not torch.cuda.is_available(): |
| raise RuntimeError("Unsloth generation requires CUDA; local Ollama is the CPU/Mac alternative") |
| adapter_path = Path(adapter).expanduser().resolve(strict=True) |
| if not (adapter_path / "adapter_config.json").exists(): |
| raise ValueError("Expected a local saved PEFT adapter") |
| config = json.loads((adapter_path / "adapter_config.json").read_text()) |
| declared_base = config["base_model_name_or_path"] |
| manifest = json.loads((adapter_path / "run_manifest.json").read_text()) |
| signature = manifest["signature"] |
| pinned_base, pinned_revision = signature["base_model"], signature["base_revision"] |
| if base_model is not None and base_model != pinned_base: |
| raise ValueError("Requested base model differs from the pinned training manifest") |
| |
| |
| self.model, self.tokenizer = FastLanguageModel.from_pretrained(model_name=pinned_base, revision=pinned_revision, |
| max_seq_length=4096, dtype=torch.float16, load_in_4bit=True, trust_remote_code=False, use_exact_model_name=True) |
| from bitsandbytes.nn import Linear4bit |
| for module in self.model.modules(): |
| if isinstance(module, Linear4bit): |
| module.compute_dtype = torch.float16 |
| module.compute_type_is_set = True |
| from peft import PeftModel |
| from transformers import AutoTokenizer |
| self.model = PeftModel.from_pretrained(self.model, str(adapter_path), is_trainable=False) |
| self.tokenizer = AutoTokenizer.from_pretrained(str(adapter_path), trust_remote_code=False) |
| FastLanguageModel.for_inference(self.model) |
| self.torch = torch |
| self.metadata = {"backend": "unsloth", "adapter": directory_fingerprint(adapter_path), "base_model": pinned_base, "adapter_declared_base": declared_base, |
| "requested_base_model": base_model, "base_revision": pinned_revision, |
| "torch_version": torch.__version__, "cuda_version": torch.version.cuda, |
| "packages": {name: importlib.metadata.version(name) for name in ["unsloth", "transformers", "peft", "bitsandbytes"]}, |
| "gpu": torch.cuda.get_device_name(), "system_prompt_sha256": hashlib.sha256(SYSTEM_PROMPT.encode()).hexdigest()} |
| try: |
| from huggingface_hub import snapshot_download |
| self.metadata["base_artifacts"] = directory_fingerprint(Path(snapshot_download(repo_id=pinned_base, revision=pinned_revision, local_files_only=True))) |
| except Exception as exc: |
| self.metadata["base_artifacts_error"] = type(exc).__name__ + ": " + str(exc) |
|
|
| def generate(self, prompt: str, max_new_tokens: int = 3072) -> dict: |
| return self.generate_many([prompt], max_new_tokens)[0] |
|
|
| def generate_many(self, prompts: list[str], max_new_tokens: int = 3072) -> list[dict]: |
| call_started = time.monotonic() |
| if not 1 <= len(prompts) <= 32: |
| raise ValueError("Free T4 inference batches must contain one to thirty-two prompts") |
| texts = [self.tokenizer.apply_chat_template( |
| [{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": prompt}], |
| tokenize=False, add_generation_prompt=True) for prompt in prompts] |
| previous_padding = self.tokenizer.padding_side |
| self.tokenizer.padding_side = "left" |
| try: |
| inputs = self.tokenizer(texts, return_tensors="pt", padding=True, add_special_tokens=False).to("cuda") |
| finally: |
| self.tokenizer.padding_side = previous_padding |
| length = inputs.input_ids.shape[1] |
| if length + max_new_tokens > 4096: |
| raise ValueError("Prompt plus generation exceeds the 4096-token bound; do not truncate evaluation prompts") |
| self.torch.cuda.synchronize() |
| start = time.monotonic() |
| with self.torch.inference_mode(): |
| output = self.model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False, |
| use_cache=True, pad_token_id=self.tokenizer.eos_token_id) |
| self.torch.cuda.synchronize() |
| elapsed = time.monotonic() - start |
| eos = self.model.generation_config.eos_token_id |
| eos_ids = set(eos if isinstance(eos, list) else [eos, self.tokenizer.eos_token_id]) |
| results = [] |
| for index, sequence in enumerate(output[:, length:].tolist()): |
| stop = next((i for i, token in enumerate(sequence) if token in eos_ids), None) |
| generated = sequence if stop is None else sequence[:stop + 1] |
| results.append({"output": self.tokenizer.decode(generated, skip_special_tokens=True), |
| "generation_seconds": elapsed, "generation_batch_size": len(prompts), |
| "latency_note": "generation_seconds times model.generate; generation_call_seconds includes chat formatting, tokenization, device transfers and decoding for the whole batch. Neither is divided by batch size.", |
| "tokens": len(generated), "prompt_tokens": int(inputs.attention_mask[index].sum().item()), |
| "done_reason": "stop" if stop is not None else "length"}) |
| call_elapsed = time.monotonic() - call_started |
| for result in results: |
| result["generation_call_seconds"] = call_elapsed |
| return results |
|
|
|
|
| def add_backend_args(parser): |
| parser.add_argument("--backend", choices=["ollama", "unsloth"], required=True) |
| parser.add_argument("--model", help="Installed Ollama model tag") |
| parser.add_argument("--ollama-url", default="http://127.0.0.1:11434") |
| parser.add_argument("--adapter", help="Local Unsloth adapter directory") |
| parser.add_argument("--base-model", default=None) |
| parser.add_argument("--max-new-tokens", type=int, default=3072) |
|
|
|
|
| def make_generator(args): |
| if not 1 <= args.max_new_tokens <= 3072: |
| raise ValueError("Generation must use 1–3072 tokens, reserving context for instructions") |
| if args.backend == "ollama": |
| if not args.model: |
| raise ValueError("--model is required for Ollama") |
| return OllamaGenerator(args.model, args.ollama_url) |
| if not args.adapter: |
| raise ValueError("--adapter is required for Unsloth") |
| return UnslothGenerator(args.adapter, args.base_model) |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser(description=__doc__) |
| add_backend_args(parser) |
| parser.add_argument("--prompt", required=True) |
| parser.add_argument("--output", type=Path, required=True) |
| args = parser.parse_args() |
| if args.output.exists(): |
| parser.error("Output already exists; refusing to overwrite a generation") |
| generator = make_generator(args) |
| result = generator.generate(args.prompt, args.max_new_tokens) |
| args.output.parent.mkdir(parents=True, exist_ok=True) |
| with args.output.open("x") as handle: |
| handle.write(result["output"]) |
| result["metadata"] = generator.metadata |
| args.output.with_suffix(args.output.suffix + ".generation.json").write_text(json.dumps(result, indent=2) + "\n") |
| parse_project(result["output"]) |
| print(args.output) |
|
|
