import streamlit as st
from llama_cpp import Llama
from huggingface_hub import hf_hub_download

st.title("🐉 Qwen GGUF Chat (CPU + llama.cpp)")

# -----------------------
# Load model (cached)
# -----------------------
@st.cache_resource
def load_model():
    model_path = hf_hub_download(
        repo_id="TheBloke/Qwen2.5-3B-Instruct-GGUF",
        filename="qwen2.5-3b-instruct.Q4_K_M.gguf"
    )

    llm = Llama(
        model_path=model_path,
        n_ctx=2048,
        n_threads=4,   # CPU threads，可调 2~8
        verbose=False
    )
    return llm

llm = load_model()

# -----------------------
# Chat history init
# -----------------------
if "messages" not in st.session_state:
    st.session_state.messages = []

# -----------------------
# Display history
# -----------------------
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# -----------------------
# Input
# -----------------------
if prompt := st.chat_input("请输入你的问题"):
    st.session_state.messages.append({"role": "user", "content": prompt})

    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):
        message_placeholder = st.empty()

        # -----------------------
        # Qwen-style prompt build
        # -----------------------
        chat_prompt = ""

        for msg in st.session_state.messages:
            if msg["role"] == "user":
                chat_prompt += f"User: {msg['content']}\n"
            else:
                chat_prompt += f"Assistant: {msg['content']}\n"

        chat_prompt += "Assistant:"

        # -----------------------
        # Generate
        # -----------------------
        output = llm(
            chat_prompt,
            max_tokens=512,      # CPU 推荐 256~512
            temperature=0.7,
            top_p=0.9,
            stop=["User:"]
        )

        response = output["choices"][0]["text"].strip()

        message_placeholder.markdown(response)

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