File size: 2,348 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 54 | 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-best"
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.2, top_p=0.9,
do_sample=True, repetition_penalty=1.05, pad_token_id=tokenizer.pad_token_id,
)
return tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
HELD_OUT = [
"Write a Python function to find the longest common prefix among a list of strings.",
"Implement a min-heap in Python from scratch without using the heapq module.",
"Write a Python function that performs topological sort on a directed acyclic graph represented as an adjacency list.",
"Write a Python function to serialize and deserialize a binary tree using a queue-based BFS approach.",
"Implement a Python function that finds all palindromic substrings of a given string.",
"Write Python code using the sliding window technique to find the longest substring without repeating characters.",
]
def main():
model, tokenizer = load_model()
for t in HELD_OUT:
print("=" * 64)
print("PROMPT:", t)
print("-" * 64)
print(generate(model, tokenizer, t))
print()
if __name__ == "__main__":
main() |