app.py

"""
Hugging Face LLM Chatbot with Gradio
Using transformers library to run models locally
"""

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

# Get HF token from environment (Spaces uses Secrets, local uses .env)
HF_TOKEN = os.getenv("HF_TOKEN", None)

# Check device
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

# Available models (optimized for local execution)
MODELS = {
    "microsoft/DialoGPT-small": {
        "name": "DialoGPT Small (영어, 빠름)",
        "max_length": 80,
        "language": "en",
    },
    "microsoft/DialoGPT-medium": {
        "name": "DialoGPT Medium (영어, 고품질)",
        "max_length": 100,
        "language": "en",
    },
    "gpt2": {
        "name": "GPT-2 (영어, 범용)",
        "max_length": 80,
        "language": "en",
    },
    "beomi/llama-2-ko-7b": {
        "name": "Llama-2-Ko 7B (한글 대화형, ⚠️ 14GB+ RAM 필요)",
        "max_length": 150,
        "language": "ko",
        "warning": "이 모델은 14GB 이상의 메모리가 필요합니다. HF Spaces 무료 tier에서는 메모리 부족으로 실행되지 않을 수 있습니다.",
    },
    "kyujinpy/KoT-Llama2-7B-Chat": {
        "name": "KoT-Llama2-7B-Chat (한글 대화, ⚠️ 14GB+ RAM 필요)",
        "max_length": 150,
        "language": "ko",
        "warning": "이 모델은 14GB 이상의 메모리가 필요합니다. HF Spaces 무료 tier에서는 메모리 부족으로 실행되지 않을 수 있습니다.",
    },
    "beomi/KoAlpaca-Polyglot-5.8B": {
        "name": "KoAlpaca 5.8B (한글 대화형, ⚠️ 12GB+ RAM 필요)",
        "max_length": 150,
        "language": "ko",
        "warning": "이 모델은 12GB 이상의 메모리가 필요합니다. HF Spaces 무료 tier에서는 메모리 부족으로 실행되지 않을 수 있습니다.",
    },
    "nlpai-lab/kullm-polyglot-5.8b-v2": {
        "name": "KULLM-Polyglot 5.8B (한글 대화, ⚠️ 12GB+ RAM 필요)",
        "max_length": 150,
        "language": "ko",
        "warning": "이 모델은 12GB 이상의 메모리가 필요합니다. HF Spaces 무료 tier에서는 메모리 부족으로 실행되지 않을 수 있습니다.",
    },
}

# Model cache
loaded_models = {}
loaded_tokenizers = {}


def load_model(model_name):
    """Load model and tokenizer"""
    if model_name not in loaded_models:
        try:
            print(f"Loading model: {model_name}")

            # Load tokenizer
            tokenizer = AutoTokenizer.from_pretrained(
                model_name,
                token=HF_TOKEN,
                padding_side='left',
                trust_remote_code=True
            )

            # Add pad token if missing
            if tokenizer.pad_token is None:
                tokenizer.pad_token = tokenizer.eos_token

            # Load model with safetensors support
            try:
                model = AutoModelForCausalLM.from_pretrained(
                    model_name,
                    token=HF_TOKEN,
                    torch_dtype=torch.float32,
                    low_cpu_mem_usage=True,
                    trust_remote_code=True,
                    use_safetensors=True
                )
            except Exception as e:
                # Fallback to default loading if safetensors fails
                print(f"⚠️ Safetensors loading failed, trying default method: {e}")
                model = AutoModelForCausalLM.from_pretrained(
                    model_name,
                    token=HF_TOKEN,
                    torch_dtype=torch.float32,
                    low_cpu_mem_usage=True,
                    trust_remote_code=True
                )

            model.to(device)
            model.eval()

            loaded_models[model_name] = model
            loaded_tokenizers[model_name] = tokenizer

            print(f"✅ Model {model_name} loaded successfully")

        except Exception as e:
            print(f"❌ Failed to load model {model_name}: {e}")
            return None, None

    return loaded_models.get(model_name), loaded_tokenizers.get(model_name)


def chat_response(message, history, model_name):
    """
    Generate chatbot response

    Args:
        message: User input
        history: Chat history in Gradio format
        model_name: Selected model

    Returns:
        Response text
    """
    try:
        # Load model and tokenizer
        model, tokenizer = load_model(model_name)

        if model is None or tokenizer is None:
            return f"❌ 모델 '{model_name}'을 로드할 수 없습니다. 다른 모델을 선택해주세요."

        model_config = MODELS[model_name]

        # Build conversation context
        conversation = ""
        for msg in history:
            if msg["role"] == "user":
                conversation += f"{msg['content']}\n"
            elif msg["role"] == "assistant":
                conversation += f"{msg['content']}\n"

        # Add current message
        conversation += f"{message}\n"

        # Tokenize
        inputs = tokenizer.encode(conversation, return_tensors="pt").to(device)

        # Generate response
        with torch.no_grad():
            outputs = model.generate(
                inputs,
                max_new_tokens=model_config["max_length"],
                temperature=0.9,
                do_sample=True,
                pad_token_id=tokenizer.pad_token_id,
                eos_token_id=tokenizer.eos_token_id,
            )

        # Decode response
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)

        # Remove the input prompt from response
        response = response[len(conversation):].strip()

        # If empty, return a default message
        if not response:
            response = "I understand. Could you tell me more?"

        return response

    except Exception as e:
        import traceback
        error_msg = str(e)
        error_type = type(e).__name__

        print("=" * 50)
        print(f"Error Type: {error_type}")
        print(f"Error Message: {error_msg}")
        print(f"Traceback:\n{traceback.format_exc()}")
        print("=" * 50)

        if "out of memory" in error_msg.lower() or "oom" in error_msg.lower():
            return "❌ 메모리 부족. 더 작은 모델을 선택하거나 앱을 재시작하세요."
        elif "cuda" in error_msg.lower() and device == "cpu":
            return "⚠️ GPU 없이 CPU로 실행 중입니다. 응답이 느릴 수 있습니다."
        else:
            return f"❌ 오류: {error_type}\n{error_msg[:200]}\n\n터미널에서 전체 로그를 확인하세요."


# Global state
current_model = "microsoft/DialoGPT-small"

# Preload default model
print("Preloading default model...")
load_model(current_model)

# Create Gradio interface
with gr.Blocks(
    title="🤖 Hugging Face Chatbot",
    theme=gr.themes.Soft(),
) as demo:
    gr.Markdown(
        """
        # 🤖 Hugging Face LLM Chatbot

        **로컬 모델 실행 방식** - API 제한 없음!

        **사용 방법:**
        1. 모델을 선택하세요 (처음에는 로딩 시간 필요)
        2. 메시지를 입력하고 대화하세요
        3. CPU에서 실행되므로 응답이 조금 느릴 수 있습니다

        **언어별 추천 모델:**
        - 🇬🇧 영어: DialoGPT, GPT-2
        - 🇰🇷 한글: KoGPT-2, KoAlpaca (5.8B는 큰 모델, 느림)

        **장점:** API 제한 없음, 완전 무료, 오프라인 작동 가능
        """
    )

    # Model selector
    model_dropdown = gr.Dropdown(
        choices=[(config["name"], model_id) for model_id, config in MODELS.items()],
        value="microsoft/DialoGPT-small",
        label="🎯 모델 선택",
        info="모델을 변경하면 새 모델을 다운로드합니다 (처음 한 번만)",
    )

    # Warning message for model requirements
    model_warning = gr.Markdown("", visible=False)

    # Chat interface
    chatbot = gr.ChatInterface(
        fn=chat_response,
        type="messages",
        additional_inputs=[model_dropdown],
        chatbot=gr.Chatbot(
            height=500,
            placeholder="메시지를 입력하세요...",
            type="messages",
        ),
        textbox=gr.Textbox(
            placeholder="메시지를 입력하세요 (영어 권장)...",
            container=False,
            scale=7,
        ),
        examples=[
            ["Hello! How are you?", "microsoft/DialoGPT-small"],
            ["Tell me a joke", "microsoft/DialoGPT-medium"],
            ["안녕하세요! 오늘 날씨가 어때요?", "beomi/llama-2-ko-7b"],
            ["인공지능에 대해 간단히 설명해주세요.", "kyujinpy/KoT-Llama2-7B-Chat"],
        ],
    )

    # Show warning and clear chat when model changes
    def on_model_change(new_model):
        global current_model
        current_model = new_model

        # Check if model has warning
        warning_text = ""
        warning_visible = False
        if "warning" in MODELS[new_model]:
            warning_text = f"⚠️ **경고**: {MODELS[new_model]['warning']}"
            warning_visible = True

        # Preload new model
        load_model(new_model)

        # Return: empty chat history, warning text, warning visibility
        return [], warning_text, gr.update(visible=warning_visible)

    model_dropdown.change(
        fn=on_model_change,
        inputs=[model_dropdown],
        outputs=[chatbot.chatbot_state, model_warning, model_warning],
    )

    gr.Markdown(
        """
        ---

        **⚠️ 참고:**
        - 모델은 로컬에서 실행됩니다 (첫 실행 시 다운로드)
        - CPU에서 실행되므로 GPU보다 느립니다
        - 각 모델은 특정 언어에 최적화되어 있습니다

        **💾 디스크 사용량:**
        - DialoGPT-small: ~350MB
        - DialoGPT-medium: ~800MB
        - GPT-2: ~500MB
        - KoGPT-2: ~500MB
        - KoAlpaca-5.8B: ~12GB (큰 모델, 메모리 8GB+ 필요)

        **💡 팁:**
        - 영어 대화는 DialoGPT 추천
        - 한글 대화는 KoGPT-2 추천 (KoAlpaca는 리소스 충분할 때만)
        - 짧은 문장으로 대화하면 더 나은 결과
        - 모델이 한 번 로드되면 다시 다운로드하지 않습니다
        """
    )

if __name__ == "__main__":
    demo.launch()