ST-Coder-14B-LoRA / interface.py
RnniaSnow's picture
Create interface.py
93799bb verified
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
# 1. Load Base Model
base_model_path = "Qwen/Qwen2.5-Coder-14B-Instruct"
model = AutoModelForCausalLM.from_pretrained(
base_model_path,
torch_dtype="auto",
device_map="auto",
trust_remote_code=True
)
tokenizer = AutoTokenizer.from_pretrained(base_model_path, trust_remote_code=True)
# 2. Load LoRA Adapter
lora_path = "RnniaSnow/ST-Coder-14B-LoRA"
model = PeftModel.from_pretrained(model, lora_path)
# 3. Generate Code
prompt = "Write a Function Block (ST) for a PID controller with anti-windup mechanism."
messages = [
{"role": "system", "content": "You are an expert IEC 61131-3 PLC programmer."},
{"role": "user", "content": prompt}
]
text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
generated_ids = model.generate(
**model_inputs,
max_new_tokens=1024,
temperature=0.2, # Low temperature for code precision
top_p=0.9
)
output_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
print(output_text)