go back again to 1b & 3b（没招了）

app.py

@@ -1,78 +1,128 @@
-import os
 import gradio as gr
-from huggingface_hub import hf_hub_download
 from llama_cpp import Llama
+from huggingface_hub import hf_hub_download
+import os
 
-# ============ 下载模型 ==============
-# 从环境变量读取 HF Token（在 Spaces → Settings → Secrets 设置）
-HF_TOKEN = os.environ.get("HF_Token")
-
-# 模型仓库与文件
-REPO_ID = "Datangtang/GGUF3B"
-FILE_NAME = "llama-3.2-3b-instruct.Q4_K_M.gguf"
-
-model_path = hf_hub_download(
-    repo_id=REPO_ID,
-    filename=FILE_NAME,
-    token=HF_TOKEN
-)
-
-# ============ 加载模型 ==============
-llm = Llama(
-    model_path=model_path,
-    n_ctx=4096,
-    n_threads=4,
-    chat_format="llama-3",
-)
-
-
-# ============ 推理函数 ==============
-def chat_fn(history, user_input):
-    """
-    history 为 Gradio 聊天历史
-    user_input 为当前用户输入
-    """
-    messages = []
-
-    # 组织对话历史，适配 llama_cpp 的聊天格式
-    for role, text in history:
-        if role == "user":
-            messages.append({"role": "user", "content": text})
-        else:
-            messages.append({"role": "assistant", "content": text})
-
-    # 新输入
-    messages.append({"role": "user", "content": user_input})
-
-    # 调用 LLM
-    result = llm.create_chat_completion(
-        messages=messages,
-        max_tokens=512,
-        temperature=0.7,
-        top_p=0.95
+# ----------------------------------------
+# Global model cache
+# ----------------------------------------
+loaded_models = {}   # Cache loaded Llama models
+current_model_name = None
+
+MODEL_CONFIGS = {
+    "1B Model (Datangtang/GGUF1B)": {
+        "repo_id": "Datangtang/GGUF1B",
+        "filename": "llama-3.2-1b-instruct.Q4_K_M.gguf"
+    },
+    "3B Model (Datangtang/GGUF3B)": {
+        "repo_id": "Datangtang/GGGF3B",
+        "filename": "llama-3.2-3b-instruct.Q4_K_M.gguf"
+    }
+}
+
+
+# ----------------------------------------
+# Load model function
+# ----------------------------------------
+def load_model(model_choice):
+    global loaded_models, current_model_name
+
+    if model_choice in loaded_models:
+        print(f"Reusing already loaded model: {model_choice}")
+        current_model_name = model_choice
+        return loaded_models[model_choice]
+
+    print(f"Downloading model: {model_choice}")
+
+    cfg = MODEL_CONFIGS[model_choice]
+
+    model_path = hf_hub_download(
+        repo_id=cfg["repo_id"],
+        filename=cfg["filename"],
+        local_dir="./model",
+        token=os.environ["HF_TOKEN"]
     )
 
-    output = result["choices"][0]["message"]["content"]
+    print(f"Model downloaded to: {model_path}")
+    print("Loading GGUF model into memory...")
+
+    llm = Llama(
+        model_path=model_path,
+        n_ctx=1024,
+        n_threads=6,
+        n_batch=512,
+        n_gpu_layers=0,
+        use_mmap=True,
+        use_mlock=True,
+        verbose=False,
+    )
+
+    loaded_models[model_choice] = llm
+    current_model_name = model_choice
+
+    print("Model loaded successfully!")
+    return llm
+
+
+# ----------------------------------------
+# Chat function
+# ----------------------------------------
+def chat(message, history, model_choice):
+    llm = load_model(model_choice)
+
+    # System prompt
+    conversation = "System: You are a helpful assistant.\n"
 
-    # 返回：更新后的历史记录
-    history.append(("user", user_input))
-    history.append(("assistant", output))
-    return history, ""
+    # Add last 3 messages
+    for human, assistant in history[-3:]:
+        conversation += f"User: {human}\nAssistant: {assistant}\n"
 
+    conversation += f"User: {message}\nAssistant:"
 
-# ============ Gradio UI ==============
+    response = llm(
+        conversation,
+        max_tokens=128,
+        temperature=0.7,
+        top_p=0.9,
+        top_k=40,
+        repeat_penalty=1.1,
+        stop=["User:", "Assistant:"],
+        echo=False
+    )
+
+    return response["choices"][0]["text"].strip()
+
+
+# ----------------------------------------
+# Gradio UI
+# ----------------------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("# 💬 Chat with Your Fine-tuned LLM")
 
-    chatbot = gr.Chatbot(height=500)
-    user_input = gr.Textbox(show_label=False, placeholder="Enter message...")
-    submit = gr.Button("Send")
+    gr.Markdown("# 🦙 Datangtang GGUF Model Demo")
+    gr.Markdown("Switch between **1B** and **3B** GGUF models in real-time.")
 
-    submit.click(
-        fn=chat_fn,
-        inputs=[chatbot, user_input],
-        outputs=[chatbot, user_input]
+    model_choice = gr.Dropdown(
+        label="Select Model",
+        choices=list(MODEL_CONFIGS.keys()),
+        value="1B Model (Datangtang/GGUF1B)",
     )
 
+    chat_iface = gr.ChatInterface(
+        fn=lambda message, history: chat(message, history, model_choice.value),
+        examples=[
+            "Explain deep learning in one paragraph.",
+            "What is the difference between supervised and unsupervised learning?",
+            "Explain what a transformer model is.",
+        ],
+        cache_examples=False,
+    )
+
+    model_choice.change(
+        fn=lambda x: f"🔄 Switched to: {x}",
+        inputs=[model_choice],
+        outputs=[],
+    )
+
+
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()