Views
No views yet
1from transformers import AutoModelForCausalLM, AutoTokenizer
2from peft import PeftModel
3import torch
4
5## Base Model
6base_model = "meta-llama/Llama-3.2-3B"
7
8## Lora
9adapter_model = "YimingZeng/FineEdit_Model"
10
11subfolder = "FineEdit-XL"
12
13## Load tokenizer and base model
14tokenizer = AutoTokenizer.from_pretrained(base_model)
15base = AutoModelForCausalLM.from_pretrained(
16 base_model,
17 torch_dtype=torch.bfloat16,
18 device_map="auto"
19)
20
21## Load LoRA adapter
22model = PeftModel.from_pretrained(base, adapter_model, subfolder=subfolder)
23
24## Test
25prompt = """Edit Request: Please change 'Captain American' to 'Iron Man'.
26Original Content: 'Captain American' is a superhero.
27Edited Content:
28"""
29inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
30
31with torch.no_grad():
32 outputs = model.generate(**inputs, max_new_tokens=100)
33
34print(tokenizer.decode(outputs[0], skip_special_tokens=True))