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