add

app.py

@@ -2,17 +2,23 @@ import streamlit as st
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# 모델 로드 (DialoGPT-medium 예시)
+# 모델 로드 (DeepSeek-R1-Distill-Qwen-1.5B 예시)
 @st.cache_resource
-def load_model(model_name="microsoft/DialoGPT-medium"):
+def load_model(model_name="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"):
     tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-    model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        device_map="auto",
+        torch_dtype=torch.float16,
+        trust_remote_code=True  # 만약 커스텀 코드가 필요한 경우 활성화
+    )
     return tokenizer, model
 
 # 앱 실행 함수
 def main():
-    st.title("ChatGPT 유사 대화 데모")
-    st.write("여기는 DialoGPT 모델을 활용한 간단한 대화 테스트용 데모입니다.")
+    st.set_page_config(page_title="DeepSeek-R1 Chatbot", page_icon="🤖")
+    st.title("DeepSeek-R1 기반 대화형 챗봇")
+    st.write("DeepSeek-R1-Distill-Qwen-1.5B 모델을 활용한 한국어 대화 테스트용 데모입니다.")
     
     # 세션 스테이트 초기화
     if "chat_history_ids" not in st.session_state:
@@ -21,9 +27,10 @@ def main():
         st.session_state["past_user_inputs"] = []
     if "generated_responses" not in st.session_state:
         st.session_state["generated_responses"] = []
-
-    tokenizer, model = load_model(model_name="deepseek-ai/DeepSeek-R1")
-
+    
+    # 모델과 토크나이저 불러오기
+    tokenizer, model = load_model()
+    
     # 기존 대화 내역 표시
     if st.session_state["past_user_inputs"]:
         for user_text, bot_text in zip(st.session_state["past_user_inputs"], st.session_state["generated_responses"]):
@@ -33,40 +40,45 @@ def main():
             # 봇 메시지
             with st.chat_message("assistant"):
                 st.write(bot_text)
-
+    
     # 채팅 입력창
     user_input = st.chat_input("메시지를 입력하세요...")
-
+    
     if user_input:
         # 사용자 메시지 표시
         with st.chat_message("user"):
             st.write(user_input)
-
-        # 새 입력을 토큰화
-        new_user_input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors='pt')
-
+        
+        # 모델 입력 전처리
+        new_user_input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors='pt').to(model.device)
+        
         if st.session_state["chat_history_ids"] is not None:
             # 기존 히스토리에 이어 붙이기
             bot_input_ids = torch.cat([st.session_state["chat_history_ids"], new_user_input_ids], dim=-1)
         else:
             bot_input_ids = new_user_input_ids
-
+        
         # 모델 추론
         with torch.no_grad():
             chat_history_ids = model.generate(
                 bot_input_ids,
-                max_length=1000,
-                pad_token_id=tokenizer.eos_token_id
+                max_length=32768,  # 모델 카드 권장 최대 길이
+                temperature=0.6,   # 모델 카드 권장 온도
+                top_p=0.95,         # 모델 카드 권장 top-p
+                pad_token_id=tokenizer.eos_token_id,
+                do_sample=True,
+                num_return_sequences=1
             )
-
-        # 결과 디코딩
-        bot_text = tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
-
+        
+        # 새로 생성된 토큰만 디코딩
+        bot_output_ids = chat_history_ids[:, bot_input_ids.shape[-1]:]
+        bot_text = tokenizer.decode(bot_output_ids[0], skip_special_tokens=True)
+        
         # 세션 스테이트에 대화 내용 업데이트
         st.session_state["past_user_inputs"].append(user_input)
         st.session_state["generated_responses"].append(bot_text)
         st.session_state["chat_history_ids"] = chat_history_ids
-
+        
         # 봇 메시지 표시
         with st.chat_message("assistant"):
             st.write(bot_text)