import os
os.environ["HF_HOME"] = "/tmp/my_unique_cache"  # avoids shared cache conflicts


os.environ["MPLCONFIGDIR"] = "/tmp"  # Prevent matplotlib config errors
os.environ["STREAMLIT_BROWSER_GATHER_USAGE_STATS"] = "false"
os.environ["STREAMLIT_SERVER_HEADLESS"] = "true"

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

# Title and UI
st.set_page_config(page_title="DeepSeek-R1 Chatbot", page_icon="🤖")
st.title("🧠 DeepSeek-R1 CPU Chatbot")
st.caption("Running entirely on CPU using Hugging Face Transformers")

# Load the model and tokenizer
@st.cache_resource
def load_model():
    tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/DeepSeek-Coder-1.3B-base")
    model = AutoModelForCausalLM.from_pretrained("deepseek-ai/DeepSeek-Coder-1.3B-base")
    return tokenizer, model

tokenizer, model = load_model()

# Prompt input
user_input = st.text_area("📥 Enter your prompt here:", "Explain what a neural network is.")

if st.button("🧠 Generate Response"):
    with st.spinner("Thinking..."):
        inputs = tokenizer(user_input, return_tensors="pt")
        outputs = model.generate(**inputs, max_new_tokens=100)
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        st.markdown("### 🤖 Response:")
        st.write(response)