Upload infer.py

infer.py

@@ -0,0 +1,176 @@
+#!/usr/bin/env python3
+"""fingpt — inference with a LoRA adapter.
+
+Loads the base model from HuggingFace Hub, injects LoRA layers using the
+metadata stored in the adapter checkpoint, then runs generation.
+
+Usage
+-----
+# Interactive REPL
+python infer.py --adapter weights_lora_coder_1b5/adapter_final.pt
+
+# Single prompt
+python infer.py --adapter weights_lora_coder_1b5/adapter_final.pt \
+    --prompt "Fix this Python code: ..."
+
+# One-liner (pipe-friendly)
+echo "Fix: def f(n): return n * f(n)" | python infer.py \
+    --adapter weights_lora_coder_1b5/adapter_final.pt
+"""
+
+import argparse
+import sys
+from pathlib import Path
+
+import torch
+
+_HERE = Path(__file__).resolve().parent
+sys.path.insert(0, str(_HERE))
+
+from fingpt.lora import inject_lora
+
+
+# ── Model loading ─────────────────────────────────────────────────────────────
+
+def load_model(adapter_path: str):
+    """Load base model + inject LoRA + load adapter weights.
+
+    All config is read from the adapter checkpoint metadata so you never
+    need to pass model name / r / alpha manually.
+    """
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+
+    ckpt = torch.load(adapter_path, map_location="cpu", weights_only=False)
+    meta       = ckpt["meta"]
+    state_dict = ckpt["state_dict"]
+
+    model_name   = meta["model_name"]
+    lora_r       = meta["lora_r"]
+    lora_alpha   = meta["lora_alpha"]
+    lora_targets = meta["lora_target_modules"]
+
+    print(f"[infer] base={model_name}  r={lora_r}  α={lora_alpha}")
+    print(f"[infer] targets={lora_targets}")
+
+    # Tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    # Base model
+    cuda_ok = torch.cuda.is_available()
+    try:
+        import accelerate  # noqa: F401
+        load_kwargs = {"device_map": "auto"} if cuda_ok else {}
+    except ImportError:
+        load_kwargs = {}
+
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        dtype=torch.bfloat16 if cuda_ok else torch.float32,
+        trust_remote_code=True,
+        **load_kwargs,
+    )
+    if not load_kwargs:
+        device = torch.device("cuda" if cuda_ok else "cpu")
+        model = model.to(device)
+
+    # Inject LoRA (dropout=0 at inference — no regularisation needed)
+    model = inject_lora(model, target_modules=lora_targets,
+                        r=lora_r, alpha=lora_alpha, dropout=0.0)
+
+    # Load trained adapter weights
+    missing, unexpected = model.load_state_dict(state_dict, strict=False)
+    lora_missing = [k for k in missing if "lora" in k]
+    if lora_missing:
+        raise ValueError(f"Missing LoRA keys: {lora_missing}")
+    print(f"[infer] Loaded {len(state_dict)} adapter tensors from {adapter_path}")
+
+    model.eval()
+    return model, tokenizer
+
+
+# ── Generation ────────────────────────────────────────────────────────────────
+
+def generate(
+    model,
+    tokenizer,
+    prompt: str,
+    max_new_tokens: int = 512,
+    temperature: float = 0.1,
+) -> str:
+    """Format prompt as ChatML and generate a response."""
+    messages = [{"role": "user", "content": prompt}]
+    text = tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+
+    device = next(model.parameters()).device
+    inputs = tokenizer(text, return_tensors="pt").to(device)
+
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            do_sample=temperature > 0,
+            temperature=temperature if temperature > 0 else 1.0,
+            pad_token_id=tokenizer.pad_token_id,
+            eos_token_id=tokenizer.eos_token_id,
+        )
+
+    new_ids = outputs[0][inputs["input_ids"].shape[1]:]
+    return tokenizer.decode(new_ids, skip_special_tokens=True)
+
+
+# ── CLI ───────────────────────────────────────────────────────────────────────
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="fingpt LoRA inference",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog=__doc__,
+    )
+    parser.add_argument("--adapter", required=True,
+                        help="Path to adapter .pt file")
+    parser.add_argument("--prompt", default=None,
+                        help="Single prompt string (omit for interactive REPL)")
+    parser.add_argument("--max-new-tokens", type=int, default=512)
+    parser.add_argument("--temperature", type=float, default=0.1,
+                        help="0 = greedy, >0 = sampling")
+    args = parser.parse_args()
+
+    model, tokenizer = load_model(args.adapter)
+
+    if args.prompt:
+        print(generate(model, tokenizer, args.prompt,
+                       args.max_new_tokens, args.temperature))
+        return
+
+    # Check stdin (pipe mode)
+    if not sys.stdin.isatty():
+        prompt = sys.stdin.read().strip()
+        if prompt:
+            print(generate(model, tokenizer, prompt,
+                           args.max_new_tokens, args.temperature))
+        return
+
+    # Interactive REPL
+    print("[infer] Interactive mode — type 'quit' or Ctrl-D to exit.\n")
+    while True:
+        try:
+            prompt = input(">>> ").strip()
+        except (EOFError, KeyboardInterrupt):
+            print()
+            break
+        if not prompt:
+            continue
+        if prompt.lower() in ("quit", "exit", "q"):
+            break
+        print()
+        print(generate(model, tokenizer, prompt,
+                       args.max_new_tokens, args.temperature))
+        print()
+
+
+if __name__ == "__main__":
+    main()