Update src/streamlit_app.py

src/streamlit_app.py

@@ -3,9 +3,10 @@ import json
 import requests
 import streamlit as st
 from datetime import datetime
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
-# ✅ 안전한 캐시 경로 설정 (최상단 필수)
+# ✅ 안전한 캐시 경로 설정
 os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf_cache"
 os.environ["HF_HOME"] = "/tmp/hf_cache"
 os.environ["HF_DATASETS_CACHE"] = "/tmp/hf_cache"
@@ -15,23 +16,25 @@ st.set_page_config(page_title="학사일정 캘린더", layout="centered")
 st.title("📅 학사일정 캘린더 + AI 요약")
 st.markdown("NEIS API에서 학사일정을 불러오고 FullCalendar로 시각화합니다.")
 
-# ✅ 디버깅 출력
-token_present = os.environ.get("HUGGINGFACE_TOKEN") is not None
-st.write("🔐 토큰 있음:", token_present)
-st.write("✅ 캐시 경로:", os.environ.get("TRANSFORMERS_CACHE"))
-
-# ✅ Gemma 모델 로딩 함수
+# ✅ 모델 로딩 함수 (skt/A.X-4.0-Light)
 @st.cache_resource
 def load_model():
     token = os.environ.get("HUGGINGFACE_TOKEN")
-    model_id = "google/gemma-2-2b-it"
+    model_id = "skt/A.X-4.0-Light"
     cache_dir = "/tmp/hf_cache"
 
     tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=token, cache_dir=cache_dir)
-    model = AutoModelForCausalLM.from_pretrained(model_id, use_auth_token=token, cache_dir=cache_dir)
-    return pipeline("text-generation", model=model, tokenizer=tokenizer)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        use_auth_token=token,
+        torch_dtype=torch.bfloat16,
+        device_map="auto",
+        cache_dir=cache_dir
+    )
+    model.eval()
+    return tokenizer, model
 
-llm = load_model()
+tokenizer, model = load_model()
 
 # ✅ 학교 정보 가져오기
 def get_school_info(region_code, school_name, api_key):
@@ -41,7 +44,7 @@ def get_school_info(region_code, school_name, api_key):
     school = data.get("schoolInfo", [{}])[1].get("row", [{}])[0]
     return school.get("SD_SCHUL_CODE"), school.get("ATPT_OFCDC_SC_CODE")
 
-# ✅ 학사일정 가져오기 (월 단위)
+# ✅ 학사일정 가져오기
 def get_schedule(region_code, school_code, year, month, api_key):
     from_ymd = f"{year}{month:02}01"
     to_ymd = f"{year}{month:02}31"
@@ -49,13 +52,13 @@ def get_schedule(region_code, school_code, year, month, api_key):
     res = requests.get(url)
     data = res.json()
     rows = data.get("SchoolSchedule", [{}])[1].get("row", [])
-    st.write("📦 불러온 일정 raw data:", rows)
     return rows
 
 # ✅ 요약 생성
 def summarize_schedule(rows, school_name, year):
     if not rows:
         return "일정이 없어 요약할 수 없습니다."
+    
     lines = []
     for row in rows:
         date = row["AA_YMD"]
@@ -63,11 +66,26 @@ def summarize_schedule(rows, school_name, year):
         event = row["EVENT_NM"]
         lines.append(f"{dt}: {event}")
     text = "\n".join(lines)
+
     prompt = f"{school_name}가 {year}년도에 가지는 학사일정은 다음과 같습니다:\n{text}\n주요 일정을 요약해주세요."
-    st.write("📤 요약에 전달된 프롬프트:", prompt)
-    result = llm([{"role": "user", "content": prompt}])
-    st.write("📥 모델 생성 결과:", result)
-    return result[0]["generated_text"].replace(prompt, "").strip()
+
+    messages = [
+        {"role": "system", "content": "당신은 학사일정을 요약해주는 AI입니다."},
+        {"role": "user", "content": prompt}
+    ]
+
+    input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(model.device)
+
+    with torch.no_grad():
+        output = model.generate(
+            input_ids,
+            max_new_tokens=256,
+            do_sample=False,
+        )
+
+    len_prompt = input_ids.shape[1]
+    response = tokenizer.decode(output[0][len_prompt:], skip_special_tokens=True).strip()
+    return response
 
 # ✅ 지역/학교/년도/월 선택 UI
 region_options = {
@@ -83,6 +101,7 @@ with st.form("query_form"):
     month = st.selectbox("월", options=list(range(1, 13)), index=6)
     submitted = st.form_submit_button("📅 학사일정 불러오기")
 
+# ✅ 제출 처리
 if submitted:
     with st.spinner("일정 불러오는 중..."):
         api_key = os.environ.get("NEIS_API_KEY", "a69e08342c8947b4a52cd72789a5ecaf")
@@ -94,7 +113,6 @@ if submitted:
             if not schedule_rows:
                 st.info("해당 조건의 학사일정이 없습니다.")
             else:
-                # ✅ 일정 출력용 FullCalendar 생성
                 events = [
                     {
                         "title": row["EVENT_NM"],
@@ -103,7 +121,6 @@ if submitted:
                     for row in schedule_rows
                     if "AA_YMD" in row and "EVENT_NM" in row
                 ]
-                st.write("📅 FullCalendar에 전달할 events:", events)
                 event_json = json.dumps(events, ensure_ascii=False)
 
                 st.components.v1.html(f"""
@@ -130,10 +147,9 @@ if submitted:
                 </html>
                 """, height=650)
 
-                # ✅ 요약 생성 버튼 추가
                 with st.expander("✨ 1년치 요약 보기", expanded=False):
                     if st.button("🤖 요약 생성하기"):
-                        with st.spinner("Gemma 모델이 요약 중..."):
+                        with st.spinner("모델이 요약 중..."):
                             summary = summarize_schedule(schedule_rows, school_name, year)
                             st.success("요약 완료!")
                             st.markdown(f"**{school_name} {year}년 {month}월 일정 요약:**\n\n{summary}")