import streamlit as st
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

# 모델 로드 (DeepSeek-R1-Distill-Qwen-1.5B 예시)
@st.cache_resource
def load_model(model_name="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"):
    pipe = pipeline(
        "text-generation",
        model=model_name,
        device_map="auto",
        torch_dtype=torch.float16,
        trust_remote_code=True,
        truncation=True,
        max_new_tokens=2048
    )
    return pipe

# 앱 실행 함수
def main():
    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:
        st.session_state["chat_history_ids"] = None
    if "past_user_inputs" not in st.session_state:
        st.session_state["past_user_inputs"] = []
    if "generated_responses" not in st.session_state:
        st.session_state["generated_responses"] = []
    
    # 모델과 토크나이저 불러오기
    pipe = 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"]):
            # 사용자 메시지
            with st.chat_message("user"):
                st.write(user_text)
            # 봇 메시지
            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)

        # 프롬프트 생성
        prompt = f"Human: {user_input}\n\nAssistant:"
        
        # 모델 생성
        response = pipe(
            prompt,
            max_new_tokens=2048,
            temperature=0.7,
            do_sample=True,
            truncation=True,
            pad_token_id=50256
        )
        bot_text = response[0]["generated_text"]
        
        # Assistant 응답만 추출 (개선된 방식)
        try:
            bot_text = bot_text.split("Assistant:")[-1].strip()
            if "</think>" in bot_text:  # 내부 사고 과정 제거
                bot_text = bot_text.split("</think>")[-1].strip()
        except:
            bot_text = "죄송합니다. 응답을 생성하는 데 문제가 발생했습니다."
        
        # 세션 스테이트에 대화 내용 업데이트
        st.session_state["past_user_inputs"].append(user_input)
        st.session_state["generated_responses"].append(bot_text)
        
        # 봇 메시지 표시
        with st.chat_message("assistant"):
            st.write(bot_text)

if __name__ == "__main__":
    main()