import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

def generate_twin_response(prompt):
    base_model_id = "Qwen/Qwen2.5-3B-Instruct"
    adapter_dir = "./digital_twin_adapters"

    print("Loading base model in bfloat16...")
    model = AutoModelForCausalLM.from_pretrained(
        base_model_id,
        torch_dtype=torch.bfloat16,
        device_map="auto",
        trust_remote_code=True
    )

    print("Loading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(base_model_id, trust_remote_code=True)

    print("Blending fine-tuned adapters with base model...")
    model = PeftModel.from_pretrained(model, adapter_dir)

    # Format user prompt to fit Qwen ChatML format
    system_prompt = "You are Jason. Speak in the first person. Relocate your raw memories instantly."
    formatted_prompt = f"<|im_start|>system\n{system_prompt}<|im_end|>\n<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
    
    print("\n--- Digital Twin Local Generation ---")
    print(f"Prompt: {prompt}")
    
    device = next(model.parameters()).device
    inputs = tokenizer(formatted_prompt, return_tensors="pt").to(device)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=256,
            temperature=0.72,
            top_p=0.9,
            do_sample=True,
            eos_token_id=tokenizer.eos_token_id
        )

    full_output = outputs[0][inputs.input_ids.shape[1]:]
    response = tokenizer.decode(full_output, skip_special_tokens=True)
    
    print(f"\nResponse: {response.strip()}")
    print("-------------------------------------")

if __name__ == "__main__":
    import sys
    query = "Tell me about yourself. What are your core memories of Bendigo?"
    if len(sys.argv) > 1:
        query = " ".join(sys.argv[1:])
    generate_twin_response(query)