#!/usr/bin/env python3
"""
Minimum-viable inference example for Fordentinc/book-builder-bookwriter-v1.

Loads Qwen 2.5 7B base + the LoRA adapter, reads a bible file from disk,
generates a chapter, writes it to stdout.

Usage:
    pip install torch transformers peft accelerate
    python inference_example.py path/to/your_bible.txt
"""
import sys, torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel

BASE = "Qwen/Qwen2.5-7B"
ADAPTER = "Fordentinc/book-builder-bookwriter-v1"

def main():
    if len(sys.argv) < 2:
        print("usage: python inference_example.py <bible.txt>", file=sys.stderr)
        sys.exit(1)
    prompt = open(sys.argv[1]).read()
    if "### Chapter" not in prompt:
        print("ERROR: bible must end with '### Chapter\\n<title>\\n\\n'", file=sys.stderr)
        sys.exit(2)

    print(f"Loading base: {BASE}", file=sys.stderr)
    tok = AutoTokenizer.from_pretrained(ADAPTER)
    base = AutoModelForCausalLM.from_pretrained(
        BASE, torch_dtype=torch.bfloat16, device_map="auto",
    )
    print(f"Loading adapter: {ADAPTER}", file=sys.stderr)
    model = PeftModel.from_pretrained(base, ADAPTER)
    model.eval()

    print(f"Prompt length: {len(prompt)} chars", file=sys.stderr)
    inputs = tok(prompt, return_tensors="pt").to(model.device)

    with torch.inference_mode():
        out = model.generate(
            **inputs,
            max_new_tokens=2048,
            do_sample=True,
            temperature=0.8,
            top_p=0.95,
            repetition_penalty=1.05,
            pad_token_id=tok.eos_token_id,
        )

    full = tok.decode(out[0], skip_special_tokens=True)
    # Print only the generated continuation (strip the bible)
    print(full[len(prompt):])

if __name__ == "__main__":
    main()