from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, GenerationConfig
from peft import LoraConfig, PeftModelForCausalLM
import torch
    
adapter_path = "GyunYeop/midm-base-GRPO-lora-tuning-KoreanCultureQA"
lora_config = LoraConfig.from_pretrained(adapter_path)
model_name = lora_config.base_model_name_or_path
tokenizer = AutoTokenizer.from_pretrained(model_name)
generation_config= GenerationConfig.from_pretrained(model_name)

bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,                       # 4‑bit 가중치
        bnb_4bit_quant_type="nf4",               # Normal‑Float 4
        bnb_4bit_use_double_quant=True,          # double‑quant
        bnb_4bit_compute_dtype=torch.bfloat16,   # Ada, Hopper, MI300 등
        llm_int8_skip_modules=["lm_head"]        # 출력층은 FP16
    )
    
base_model = AutoModelForCausalLM.from_pretrained(model_name, 
                                                device_map="cuda:0",
                                                trust_remote_code=True,
                                                quantization_config=bnb_config,
                                                )

model = PeftModelForCausalLM.from_pretrained(base_model, adapter_path)
model.load_adapter(adapter_path, subfolder="selection", adapter_name="selection")

prompt = '단답형 문제에서 정답을 맞추기위해 반드시 충분히 생각해보고 "생각 과정:" 이후에 정답 근거 및 생각과정을 작성 한 다음 이후 최종 답변을 생성할 것 (문제 유형: 단답형 생각과정: {생각과정} 답변: {최종답변} 형태)'
question = '문제 유형: 단답형 \n 질문: 2005년에 개관하였으며, 교육 및 문화적 목적으로 영화를 상영하는 서울의 유일한 비영리 민간 시네마테크 전용관은 어디인가요?'
inputs = tokenizer.apply_chat_template(
    [
        {"role": "system", "content": prompt},
        {"role": "user", "content": question}
    ],
    tokenize=True,
    add_generation_prompt=True,
    return_tensors="pt"
)

answer = model.generate(
    inputs.to("cuda:0"),
    generation_config=generation_config,
    max_new_tokens=1024,
    do_sample=True
)

answer_text = tokenizer.decode(answer[0][inputs.shape[-1]:], skip_special_tokens=True)
print(answer_text)