from transformers import pipeline
import gradio as gr

# KoAlpaca 모델 로드 (속도와 품질의 밸런스)
generator = pipeline(
    "text-generation",
    model="beomi/KoAlpaca-Polyglot-1.1B",
    tokenizer="beomi/KoAlpaca-Polyglot-1.1B",
    device_map="auto"
)

def answer_question(prompt):
    system_prompt = (
        "너는 한국 대학 입시 정보를 알려주는 AI야. "
        "수능, 학생부종합, 논술, 정시 등 전형을 이해하기 쉽게 설명해줘.\n\n"
    )
    response = generator(
        system_prompt + prompt,
        max_new_tokens=250,
        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="예: 가천대 논술전형 / 2025 수능 일정 / 한국공학대 전형요약"
    ),
    outputs=gr.Textbox(label="AI 답변"),
    title="입시 전문 AI (KoAlpaca 1.1B)",
    description="빠르고 똑똑한 입시정보 챗봇입니다. 실제 전형정보에 근거한 설명을 제공합니다."
)

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