from transformers import AutoModelForCausalLM, AutoTokenizer
import gradio as gr
import torch

# CPU 강제 설정
device = "cpu"
torch.set_num_threads(4)  # CPU 스레드 제한

# 경량화된 모델 로드
model = AutoModelForCausalLM.from_pretrained(
    "naver-hyperclovax/HyperCLOVAX-SEED-Text-Instruct-0.5B",
    torch_dtype=torch.float32,  # CPU는 float32 권장
    low_cpu_mem_usage=True
).to(device)

tokenizer = AutoTokenizer.from_pretrained(
    "naver-hyperclovax/HyperCLOVAX-SEED-Text-Instruct-0.5B"
)

def predict(message, history):
    # 메모리 절약을 위한 간소화된 채팅 구성
    chat = [
        {"role": "system", "content": "간결하게 답변해주세요."},
        {"role": "user", "content": message}
    ]
    
    # CPU 최적화 설정
    inputs = tokenizer.apply_chat_template(
        chat,
        return_tensors="pt",
        max_length=512,  # 길이 제한
        truncation=True
    ).to(device)
    
    outputs = model.generate(
        inputs,
        max_new_tokens=200,  # 짧은 응답
        temperature=0.3,    # 창의성 감소
        do_sample=False      #确定性 응답
    )
    
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# 경량화된 인터페이스
demo = gr.ChatInterface(
    predict,
    title="CLOVA X (CPU 모드)",
    description="CPU 전용 경량화 버전",
    theme="soft",
    examples=["안녕하세요", "날씨 알려줘"]
)

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        favicon_path=None,
        prevent_thread_lock=True
    )