from transformers import pipeline
import gradio as gr

# 빠른 한국어 GPT-2 모델 로드
generator = pipeline(
    "text-generation",
    model="skt/kogpt2-base-v2",
    tokenizer="skt/kogpt2-base-v2"
)

def answer_question(prompt):
    system_prompt = (
        "너는 한국 입시정보를 알려주는 AI야. "
        "질문에 맞게 간단명료하게 설명해줘.\n\n"
    )
    response = generator(
        system_prompt + prompt,
        max_new_tokens=150,
        temperature=0.7,
        top_p=0.9,
        do_sample=True
    )
    return response[0]["generated_text"].replace(system_prompt, "").strip()

app = gr.Interface(
    fn=answer_question,
    inputs=gr.Textbox(
        lines=2,
        label="입시 질문 입력",
        placeholder="예: 수능 일정 / 한국공학대 논술전형 / 가천대 모집인원"
    ),
    outputs=gr.Textbox(label="AI 답변"),
    title="입시정보 AI (빠른 버전)",
    description="CPU에서도 빠르게 작동하는 입시전문 챗봇입니다."
)

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