File size: 2,173 Bytes
f238825 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import PeftModel
BASE_MODEL = "Qwen/Qwen2.5-Coder-7B-Instruct"
ADAPTER_PATH = "/home/ai/qwen-coder-programming-finetuned"
def load_model():
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_use_double_quant=True,
)
model = AutoModelForCausalLM.from_pretrained(
BASE_MODEL, quantization_config=bnb_config, device_map="auto",
trust_remote_code=True, torch_dtype=torch.bfloat16,
)
model = PeftModel.from_pretrained(model, ADAPTER_PATH)
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True, use_fast=True)
tokenizer.pad_token = tokenizer.eos_token
return model, tokenizer
def generate(model, tokenizer, prompt):
messages = [{"role": "user", "content": prompt}]
text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(text, return_tensors="pt").to(model.device)
with torch.no_grad():
outputs = model.generate(
**inputs, max_new_tokens=512, temperature=0.3, top_p=0.9,
do_sample=True, repetition_penalty=1.1, pad_token_id=tokenizer.pad_token_id,
)
return tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
def main():
model, tokenizer = load_model()
tests = [
"Write a Python function to merge two sorted lists into one sorted list.",
"Write a Python class implementing a simple LRU cache with get and put operations in O(1).",
"Write a Python function using dynamic programming to compute the length of the longest increasing subsequence in an array.",
"Write Python code using Dijkstra's algorithm to find the shortest path in a weighted graph.",
]
for t in tests:
print("=" * 60)
print("PROMPT:", t)
print("-" * 60)
print(generate(model, tokenizer, t))
print()
if __name__ == "__main__":
main() |