57 lines
1.8 KiB
Python
57 lines
1.8 KiB
Python
|
#!/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()
|