Update app.py

app.py

@@ -1,32 +1,18 @@
+from transformers import AutoTokenizer, AutoModelForMaskedLM, pipeline
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForQuestionAnswering, pipeline
 
-# 모델과 토크나이저 불러오기
-model_name = "nlpai-lab/kobert-base-v1"
+model_name = "kykim/bert-kor-base"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForQuestionAnswering.from_pretrained(model_name)
+model = AutoModelForMaskedLM.from_pretrained(model_name)
+fill_mask = pipeline("fill-mask", model=model, tokenizer=tokenizer)
 
-# 질문응답 파이프라인 생성
-qa_pipeline = pipeline("question-answering", model=model, tokenizer=tokenizer)
+def answer(text):
+    masked = text.replace("___", fill_mask.tokenizer.mask_token)
+    result = fill_mask(masked)
+    return result[0]["sequence"]
 
-# 임시 예시 문맥 (나중에 입시 정보로 교체 가능)
-context = """
-한국장학재단은 국가장학금과 학자금대출 등을 관리하는 기관입니다.
-수능은 매년 11월 셋째 주 목요일에 시행됩니다.
-정시모집은 수능 성적 위주로 선발하며, 수시모집은 학생부나 면접 등 다양한 전형을 포함합니다.
-"""
-
-def answer(question):
-    try:
-        result = qa_pipeline(question=question, context=context)
-        return result["answer"]
-    except:
-        return "죄송합니다. 그 질문에 대한 답을 찾지 못했어요."
-
-gr.Interface(
-    fn=answer,
-    inputs="text",
-    outputs="text",
-    title="입시 도우미 AI (KoBERT)",
-    description="한국어 KoBERT 모델을 사용한 입시 정보 Q&A AI입니다."
-).launch()
+gr.Interface(fn=answer,
+             inputs="text",
+             outputs="text",
+             title="입시 AI (BERT 경량 버전)",
+             description="문장 완성 기반 입시 정보 챗봇").launch()