app.py

"""
무료 AI 챗봇 - HuggingFace Spaces
간단하고 안정적인 버전
"""

import gradio as gr
from transformers import pipeline
import os

# 환경 확인
IS_SPACES = os.getenv("SPACE_ID") is not None

print(f"Starting in {'HuggingFace Spaces' if IS_SPACES else 'Local'} mode...")

# 모델 로드
print("Loading model...")
generator = pipeline(
    'text-generation',
    model='microsoft/DialoGPT-small',
    device=-1  # CPU
)
print("Model loaded!")

def chat(message, history):
    """챗봇 응답 생성"""
    if not message:
        return "메시지를 입력해주세요."
    
    # 히스토리가 없으면 빈 리스트로
    if history is None:
        history = []
    
    # 간단한 프롬프트
    prompt = f"User: {message}\nBot:"
    
    try:
        # 텍스트 생성
        response = generator(
            prompt,
            max_new_tokens=50,
            temperature=0.7,
            do_sample=True,
            pad_token_id=50256
        )
        
        # 응답 추출
        generated = response[0]['generated_text']
        
        # Bot: 이후 텍스트만 추출
        if "Bot:" in generated:
            reply = generated.split("Bot:")[-1].strip()
        else:
            reply = "안녕하세요! 무엇을 도와드릴까요?"
        
        return reply
        
    except Exception as e:
        print(f"Error: {e}")
        return "죄송합니다. 오류가 발생했습니다."

# Gradio 인터페이스 (최대한 간단하게)
demo = gr.ChatInterface(
    fn=chat,
    title="무료 AI 챗봇",
    description="DialoGPT 모델을 사용한 간단한 챗봇",
    examples=["안녕하세요", "오늘 날씨 어때요?", "AI가 뭔가요?"],
)

if __name__ == "__main__":
    # HuggingFace Spaces에서 실행시 설정
    if IS_SPACES:
        demo.launch(
            server_name="0.0.0.0",
            server_port=7860
        )
    else:
        # 로컬에서 실행시
        demo.launch()