import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import PeftModel

# 모델 설정 (여기를 수정하세요!)
MODELS = {
    # ========================================
    # 03번: 한국어 요약 (EXAONE-3.5)
    # ========================================
    "한국어 요약 (03번)": {
        "base_model": "LGAI-EXAONE/EXAONE-3.5-2.4B-Instruct",
        "lora_path": "your-username/exaone-summary-lora",  # TODO: 본인 경로로!
        "prompt_template": "{input}\n\n요약:",
        "max_new_tokens": 60,
        "placeholder": "뉴스 기사를 입력하세요...",
        "example": "서울시가 내년부터 전기차 충전소를 대폭 확대한다.",
    },
    
    # ========================================
    # 05번: 다중 모델 비교 (선택사항)
    # ========================================
    # "Granite 요약 (05번)": {
    #     "base_model": "ibm-granite/granite-4.0-micro",
    #     "lora_path": "your-username/granite-summary-lora",
    #     "prompt_template": "<|user|>\n{input}\n\n위 기사를 요약해주세요.<|assistant|>\n",
    #     "max_new_tokens": 60,
    #     "placeholder": "뉴스 기사를 입력하세요...",
    #     "example": "서울시가 내년부터 전기차 충전소를 대폭 확대한다.",
    # },
    # "Qwen3 요약 (05번)": {
    #     "base_model": "Qwen/Qwen3-4B-Instruct-2507",
    #     "lora_path": "your-username/qwen3-summary-lora",
    #     "prompt_template": "<|im_start|>user\n{input}\n\n위 기사를 요약해주세요.<|im_end|>\n<|im_start|>assistant\n",
    #     "max_new_tokens": 60,
    #     "placeholder": "뉴스 기사를 입력하세요...",
    #     "example": "서울시가 내년부터 전기차 충전소를 대폭 확대한다.",
    # },
    
    # ========================================
    # 06번: 감정 분류 + 영어 QA
    # ========================================
    "감정 분류 (06번)": {
        "base_model": "LGAI-EXAONE/EXAONE-3.5-2.4B-Instruct",
        "lora_path": "your-username/lora-sentiment",
        "prompt_template": "다음 영화 리뷰의 감정을 분류하세요.\n\n리뷰: {input}\n\n감정:",
        "max_new_tokens": 10,
        "placeholder": "영화 리뷰를 입력하세요...",
        "example": "This movie was amazing! Great story and excellent acting.",
    },
    "영어 QA (06번)": {
        "base_model": "LGAI-EXAONE/EXAONE-3.5-2.4B-Instruct",
        "lora_path": "your-username/lora-qa",
        "prompt_template": "Context: The Eiffel Tower is in Paris, France.\n\nQuestion: {input}\n\nAnswer:",
        "max_new_tokens": 30,
        "placeholder": "질문을 입력하세요...",
        "example": "Where is the Eiffel Tower located?",
    },
}

loaded_models = {}

def load_model(model_name):
    if model_name in loaded_models:
        return loaded_models[model_name]
    
    config = MODELS[model_name]
    
    tokenizer = AutoTokenizer.from_pretrained(config["base_model"], use_fast=False)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    
    quant_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_use_double_quant=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16,
    )
    
    base_model = AutoModelForCausalLM.from_pretrained(
        config["base_model"],
        device_map="auto",
        trust_remote_code=True,
        quantization_config=quant_config,
    )
    
    model = PeftModel.from_pretrained(base_model, config["lora_path"])
    
    loaded_models[model_name] = (model, tokenizer, config)
    return model, tokenizer, config

def generate_response(model_name, user_input):
    try:
        model, tokenizer, config = load_model(model_name)
        prompt = config["prompt_template"].format(input=user_input)
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=config["max_new_tokens"],
                temperature=0.7,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id,
            )
        
        result = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # 프롬프트 제거
        if "요약:" in result:
            return result.split("요약:")[-1].strip()
        elif "감정:" in result:
            return result.split("감정:")[-1].strip()
        elif "Answer:" in result:
            return result.split("Answer:")[-1].strip()
        elif "<|assistant|>" in result:
            return result.split("<|assistant|>")[-1].strip()
        elif "<|im_start|>assistant" in result:
            return result.split("<|im_start|>assistant")[-1].replace("<|im_end|>", "").strip()
        else:
            return result[len(prompt):].strip()
    except Exception as e:
        return f"❌ 오류: {str(e)}"

with gr.Blocks(title="LoRA 모델 데모") as demo:
    gr.Markdown("# 🤖 LoRA 파인튜닝 모델 데모")
    gr.Markdown("Day 1에서 학습한 여러 LoRA 모델을 테스트해보세요!")
    
    with gr.Row():
        with gr.Column():
            model_dropdown = gr.Dropdown(
                choices=list(MODELS.keys()),
                value=list(MODELS.keys())[0],
                label="📌 모델 선택"
            )
            input_text = gr.Textbox(label="💬 입력", lines=5)
            submit_btn = gr.Button("🚀 실행", variant="primary")
        
        with gr.Column():
            output_text = gr.Textbox(label="✨ 결과", lines=10)
    
    submit_btn.click(
        fn=generate_response,
        inputs=[model_dropdown, input_text],
        outputs=[output_text]
    )

demo.launch()