app.py

# app.py
import gradio as gr
import requests
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM

# Load KoAlpaca model
model_id = "beomi/KoAlpaca-Polyglot-5.8B"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
generator = pipeline("text-generation", model=model, tokenizer=tokenizer)

NEIS_KEY = "a69e08342c8947b4a52cd72789a5ecaf"
SCHOOL_INFO_URL = "https://open.neis.go.kr/hub/schoolInfo"
SCHEDULE_URL = "https://open.neis.go.kr/hub/SchoolSchedule"

REGIONS = {
    "서울특별시교육청": "B10",
    "경상북도교육청": "R10"
}

MONTH_NAMES = ["01", "02", "03", "04", "05", "06", "07", "08", "09", "10", "11", "12"]


def get_school_code(region_code, school_name):
    params = {
        "KEY": NEIS_KEY,
        "Type": "json",
        "pIndex": 1,
        "pSize": 1,
        "SCHUL_NM": school_name,
        "ATPT_OFCDC_SC_CODE": region_code
    }
    res = requests.get(SCHOOL_INFO_URL, params=params)
    data = res.json()
    try:
        return data["schoolInfo"][1]["row"][0]["SD_SCHUL_CODE"], data["schoolInfo"][1]["row"][0]["ATPT_OFCDC_SC_CODE"]
    except:
        return None, None


def get_schedule(region_code, school_code, year, month):
    from_ymd = f"{year}{month}01"
    to_ymd = f"{year}{month}31"
    params = {
        "KEY": NEIS_KEY,
        "Type": "json",
        "pIndex": 1,
        "pSize": 100,
        "ATPT_OFCDC_SC_CODE": region_code,
        "SD_SCHUL_CODE": school_code,
        "AA_FROM_YMD": from_ymd,
        "AA_TO_YMD": to_ymd
    }
    res = requests.get(SCHEDULE_URL, params=params)
    data = res.json()
    try:
        rows = data["SchoolSchedule"][1]["row"]
        return rows
    except:
        return []


def generate_answer(region, school_name, year, month, question):
    region_code = REGIONS.get(region)
    if not region_code:
        return "잘못된 교육청입니다."

    school_code, confirmed_region = get_school_code(region_code, school_name)
    if not school_code:
        return "학교 정보를 찾을 수 없습니다."

    schedule_rows = get_schedule(confirmed_region, school_code, year, month)
    if not schedule_rows:
        schedule_text = "현재 일정 정보는 없습니다."
    else:
        schedule_text = "\n".join(f"{row['AA_YMD']}: {row['EVENT_NM']}" for row in schedule_rows)

    prompt = f"""일정 정보:
{schedule_text}

사용자 질문: {question}

자연스럽게 대답하세요."""

    result = generator(prompt, max_new_tokens=200, temperature=0.7)[0]["generated_text"]
    return result


def interface_fn(region, school_name, year, month, question):
    return generate_answer(region, school_name, year, month, question)


with gr.Interface(
    fn=interface_fn,
    inputs=[
        gr.Dropdown(choices=list(REGIONS.keys()), label="교육청 선택"),
        gr.Textbox(label="학교명 입력"),
        gr.Textbox(label="년도 입력", placeholder="예: 2025"),
        gr.Dropdown(choices=MONTH_NAMES, label="월 선택 (예: 07)"),
        gr.Textbox(label="GPT 질문 입력")
    ],
    outputs=gr.Textbox(label="GPT의 응답"),
    title="학사일정 + GPT 챗봇 (KoAlpaca)"
) as app:
    app.launch()