Update app.py

app.py

@@ -2,70 +2,68 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 import torch
 from threading import Thread
+import os  # 1. Secret을 읽기 위해 os 모듈 임포트
 
 # --- 1. 모델 로드 (Space의 GPU 활용) ---
-# 2505 모델은 아직 존재하지 않아, 현재 최신 모델인 2405로 수정했습니다.
 MODEL_NAME = "kakaocorp/kanana-1.5-2.1b-instruct-2405" 
 
+# 2. Space Setting에 등록한 Secret(HF_TOKEN)을 환경 변수에서 가져옵니다.
+HF_AUTH_TOKEN = os.environ.get("HF_TOKEN")
+
 print(f"모델을 로딩 중입니다: {MODEL_NAME} (Space GPU 사용)")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+
+# 3. 토큰을 사용하여 인증된 상태로 모델을 다운로드합니다.
+tokenizer = AutoTokenizer.from_pretrained(
+    MODEL_NAME,
+    token=HF_AUTH_TOKEN  # 인증 토큰 전달
+)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_NAME, 
-    torch_dtype=torch.float16,  # 메모리 절약을 위해 16비트 사용
-    device_map="auto"           # 중요: 알아서 GPU에 할당
+    torch_dtype=torch.float16,
+    device_map="auto",
+    token=HF_AUTH_TOKEN  # 인증 토큰 전달
 )
 print("모델 로딩 완료!")
 
 # --- 2. 챗봇 응답 함수 (Gradio가 이 함수를 호출) ---
-# message: 사용자가 입력한 메시지
-# history: 이전 대화 기록 (Gradio가 자동으로 관리)
 def predict(message, history):
     
-    # Kanana의 프롬프트 형식: <bos>user\n{prompt}\n<eos>assistant\n
     history_prompt = ""
-    # 이전 대화 기록(history)을 Kanana 프롬프트 형식으로 변환
     for user_msg, assistant_msg in history:
         history_prompt += f"<bos>user\n{user_msg}\n<eos>assistant\n{assistant_msg}\n"
     
-    # 현재 메시지를 프롬프트에 추가
     final_prompt = history_prompt + f"<bos>user\n{message}\n<eos>assistant\n"
     
     inputs = tokenizer(final_prompt, return_tensors="pt").to(model.device)
     
-    # --- 실시간 타이핑 효과(스트리밍)를 위한 설정 ---
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
     
-    # 모델 생성(generate) 작업을 별도 스레드에서 실행
     generation_kwargs = dict(
-        **inputs, # inputs 딕셔너리의 모든 키-값 쌍을 인자로 전달
+        **inputs,
         streamer=streamer,
-        max_new_tokens=1024, # 최대 생성 토큰 수
+        max_new_tokens=1024,
         eos_token_id=tokenizer.eos_token_id,
         pad_token_id=tokenizer.pad_token_id,
-        temperature=0.7,   # 창의성 조절
-        do_sample=True     # 샘플링 사용
+        temperature=0.7,
+        do_sample=True
     )
     
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
     
-    # 스트리머에서 나오는 텍스트를 바로바로 반환 (yield)
     generated_text = ""
     for new_text in streamer:
         generated_text += new_text
-        yield generated_text # 텍스트를 한 글자씩 실시간으로 보냄
+        yield generated_text
 
 # --- 3. Gradio 챗봇 UI 생성 ---
-# gr.ChatInterface를 쓰면 UI가 챗봇 형태로 자동 생성됩니다.
 chatbot_ui = gr.ChatInterface(
-    fn=predict,  # 챗봇이 사용할 함수
+    fn=predict,
     title="Kanana 1.5 챗봇 테스트 🤖",
     description=f"{MODEL_NAME} 모델을 테스트합니다.",
-    theme="soft", # 테마 설정
+    theme="soft",
     examples=[["한국의 수도는 어디야?"], ["AI에 대해 3줄로 요약해줘."]]
-    # retry_btn, undo_btn, clear_btn 파라미터는 현재 Gradio 버전에서 지원되지 않아 삭제했습니다.
 )
 
 # --- 4. 앱 실행 ---
-# .launch()로 Space에서 앱을 실행시킵니다.
 chatbot_ui.launch()