Views
No views yet
1 import torch
2 from transformers import AutoModelForCausalLM, AutoTokenizer
3 from peft import PeftModel
4
5 BASE = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
6 ADAPTER = "./peft" # or "org/repo-name" if pushed to HF
7
8 # Tokenizer (includes the chat template)
9 tokenizer = AutoTokenizer.from_pretrained(BASE)
10
11 # Base model (GPU, 8-bit). For CPU, remove load_in_8bit and device_map.
12 model = AutoModelForCausalLM.from_pretrained(
13 BASE,
14 device_map="auto",
15 load_in_8bit=True,
16 )
17 model = PeftModel.from_pretrained(model, ADAPTER)
18 model.eval()
19
20 # Chat prompt via tokenizer's chat_template
21 messages = [
22 {"role": "system", "content": "You are a helpful assistant for sanctions/AML."},
23 {"role": "user", "content": "Summarize the key OFAC FAQ topics."},
24 ]
25 inputs = tokenizer.apply_chat_template(
26 messages, add_generation_prompt=True, return_tensors="pt"
27 ).to(model.device)
28
29 with torch.inference_mode():
30 out = model.generate(
31 inputs,
32 max_new_tokens=256,
33 temperature=0.7,
34 top_p=0.9,
35 do_sample=True,
36 pad_token_id=tokenizer.eos_token_id,
37 )
38
39 print(tokenizer.decode(out[0], skip_special_tokens=True))
40