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
from dotenv import load_dotenv

# Load environment variables
load_dotenv()
HF_TOKEN = os.getenv("HF_TOKEN")

# 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",
    },
    "skt/kogpt2-base-v2": {
        "name": "KoGPT-2 (한글 특화)",
        "max_length": 100,
        "language": "ko",
    },
    "beomi/KoAlpaca-Polyglot-5.8B": {
        "name": "KoAlpaca 5.8B (한글 대화형, 느림)",
        "max_length": 150,
        "language": "ko",
    },
}

# 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'
            )

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

            # Load model
            model = AutoModelForCausalLM.from_pretrained(
                model_name,
                token=HF_TOKEN,
                torch_dtype=torch.float32,
            )
            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="모델을 변경하면 새 모델을 다운로드합니다 (처음 한 번만)",
    )

    # 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"],
            ["안녕하세요! 오늘 날씨가 좋네요.", "skt/kogpt2-base-v2"],
            ["인공지능에 대해 설명해주세요.", "skt/kogpt2-base-v2"],
        ],
    )

    # Clear chat when model changes
    def on_model_change(new_model):
        global current_model
        current_model = new_model
        # Preload new model
        load_model(new_model)
        return None

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

    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()