from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# SOLAR 10.7B Instruct 모델 로드
model_name = "upstage/SOLAR-10.7B-Instruct-v1.0"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto",
)

# 대화 형식의 프롬프트 구성
conversation = [
    {"role": "system", "content": "당신은 한국어 AI 어시스턴트입니다. 친절하고 정확하게 답변하세요."},
    {"role": "user", "content": "사전학습 언어 모델이란 무엇인가요? 간단히 설명해주세요."},
]

# 토크나이저의 chat_template 적용
prompt = tokenizer.apply_chat_template(
    conversation,
    tokenize=False,
    add_generation_prompt=True,
)

inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

with torch.no_grad():
    outputs = model.generate(
        **inputs,
        max_new_tokens=256,
        temperature=0.7,
        top_p=0.9,
        do_sample=True,
        repetition_penalty=1.1,
    )

# 입력 부분을 제외한 생성된 텍스트만 출력
generated = outputs[0][inputs["input_ids"].shape[-1]:]
response = tokenizer.decode(generated, skip_special_tokens=True)
print(response)

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# EXAONE 3.5 7.8B Instruct 모델 로드
model_name = "LGAI-EXAONE/EXAONE-3.5-7.8B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    trust_remote_code=True,
)

# EXAONE 대화 형식
messages = [
    {"role": "system", "content": "당신은 도움이 되는 AI 어시스턴트입니다."},
    {"role": "user", "content": "Transformer 모델의 핵심 구성 요소 3가지를 설명해주세요."},
]

# 프롬프트 생성 및 추론
input_ids = tokenizer.apply_chat_template(
    messages,
    tokenize=True,
    add_generation_prompt=True,
    return_tensors="pt",
).to(model.device)

with torch.no_grad():
    output = model.generate(
        input_ids,
        max_new_tokens=512,
        temperature=0.7,
        top_p=0.9,
        do_sample=True,
    )

response = tokenizer.decode(output[0][input_ids.shape[-1]:], skip_special_tokens=True)
print(response)

from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
import torch

# 4-bit 양자화 설정 (GPU 메모리 약 75% 절약)
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
)

# SOLAR 10.7B를 4-bit로 로드 (약 6GB VRAM)
model_name = "upstage/SOLAR-10.7B-Instruct-v1.0"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map="auto",
)

print(f"모델 메모리 사용량: {model.get_memory_footprint() / 1e9:.1f} GB")