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()