import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

# Model ve tokenizer yüklemeleri için quantize ayarları
bnb_config = BitsAndBytesConfig(
    load_in_8bit=True,
    llm_int8_threshold=6.0
)

@st.cache_resource
def load_model():
    try:
        tokenizer = AutoTokenizer.from_pretrained(
            "deepseek-ai/deepseek-r1",
            trust_remote_code=True
        )
        
        model = AutoModelForCausalLM.from_pretrained(
            "deepseek-ai/deepseek-r1",
            trust_remote_code=True,
            quantization_config=bnb_config,
            device_map="auto",
            low_cpu_mem_usage=True
        )
        return tokenizer, model
    except Exception as e:
        st.error(f"Model yüklenemedi: {str(e)}")
        st.stop()

tokenizer, model = load_model()

# Streamlit arayüzü
st.title("DeepSeek-R1 ChatBot")
st.markdown("""
**Not:** Bu demo modeli CPU üzerinde çalışmaktadır. Yanıt süreleri GPU'ya göre daha uzun olabilir.
""")

# Sohbet geçmişi yönetimi
if "messages" not in st.session_state:
    st.session_state.messages = []

# Geçmiş mesajları göster
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# Kullanıcı girişi
prompt = st.chat_input("Mesajınızı yazın...")

if prompt:
    # Kullanıcı mesajını ekle
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    # Asistan yanıtı oluştur
    with st.chat_message("assistant"):
        message_placeholder = st.empty()
        full_response = ""
        
        try:
            # Model için prompt formatı
            inputs = tokenizer.apply_chat_template(
                [{"role": "user", "content": prompt}],
                return_tensors="pt"
            )
            
            # Generation ayarları
            outputs = model.generate(
                inputs.to(model.device),
                max_new_tokens=512,
                temperature=0.7,
                top_p=0.9,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id
            )
            
            full_response = tokenizer.decode(
                outputs[0][len(inputs[0]):], 
                skip_special_tokens=True
            )
            
        except Exception as e:
            full_response = f"Hata oluştu: {str(e)}"
        
        message_placeholder.markdown(full_response)

    st.session_state.messages.append(
        {"role": "assistant", "content": full_response}
    )