import gradio as gr
from llama_cpp import Llama
from huggingface_hub import hf_hub_download

# 下载 GGUF 模型
model_path = hf_hub_download(
    repo_id="KANGYONGMA/Chemistry",
    filename="Phi-3mini.gguf"
)

# 加载模型（CPU only）
llm = Llama(
    model_path=model_path,
    n_threads=4,       # 根据 Space CPU 配置调整
    n_ctx=2048,        # 上下文长度
    verbose=False
)

SYSTEM_PROMPT = (
    "You are GVIM, an AI chemistry assistant. "
    "Your role is to help users solve problems in chemistry, "
    "including chemical reactions, molecular structures, data analysis, "
    "and providing clear scientific explanations."
)

def chat(prompt):
    full_prompt = f"{SYSTEM_PROMPT}\nUser: {prompt}\nAssistant:"
    output = llm(full_prompt, max_tokens=300, stop=["User:", "Assistant:"])
    return output["choices"][0]["text"].strip()

demo = gr.Interface(
    fn=chat,
    inputs=gr.Textbox(lines=5, placeholder="Ask your chemistry question here..."),
    outputs="text",
    title="GVIM Chemistry Assistant",
    description="An AI assistant specialized in chemistry. Ask me about reactions, molecular design, or research help."
)

demo.launch()