Switch to Hugging Face Inference API for efficient model access

app.py

@@ -1,160 +1,127 @@
 import gradio as gr
-import torch
 import os
-import gc
-from datetime import datetime
+from huggingface_hub import InferenceClient
 
-# Patch for is_torch_fx_available if needed
-try:
-        from transformers.utils import is_torch_fx_available
-except ImportError:
-        import transformers.utils
-        transformers.utils.is_torch_fx_available = lambda: False
-        print("Patched is_torch_fx_available function")
-
-    from transformers import AutoTokenizer, AutoModelForCausalLM
-
-# Model configuration
+# Model configuration - Using Inference API
 MODEL_NAME = "moonshotai/Kimi-K2-Instruct"
-DEFAULT_SYSTEM_PROMPT = "You are Kimi, an AI assistant created by Moonshot AI."
-
-# Global variables
-model = None
-tokenizer = None
-
-def load_model():
-        """Load the Kimi-K2-Instruct model with optimized settings"""""
-        global model, tokenizer
-
-        print("=" * 50)
-        print(f"Starting at {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
-        print("=" * 50)
-        print(f"CUDA available: {torch.cuda.is_available()}")
-
-        if torch.cuda.is_available():
-                    print(f"GPU count: {torch.cuda.device_count()}")
-                    for i in range(torch.cuda.device_count()):
-                                    gpu_name = torch.cuda.get_device_name(i)
-                                    gpu_mem = torch.cuda.get_device_properties(i).total_memory / 1e9
-                                    print(f"  GPU {i}: {gpu_name} ({gpu_mem:.1f} GB)")
-
-                    # Clear GPU memory
-                    torch.cuda.empty_cache()
-                    gc.collect()
-
-        print(f"\nLoading model: {MODEL_NAME}")
-
-        try:
-                    # Load tokenizer first
-                    print("Loading tokenizer...")
-                    tokenizer = AutoTokenizer.from_pretrained(
-                                    MODEL_NAME,
-                                    trust_remote_code=True,
-                                    token=os.environ.get("HF_TOKEN")
-                    )
-                    print("Tokenizer loaded successfully")
-
-                    # Load model with memory optimizations
-                    print("Loading model (this may take several minutes)...")
-                    model = AutoModelForCausalLM.from_pretrained(
-                                    MODEL_NAME,
-                                    torch_dtype=torch.bfloat16,
-                                    device_map="auto",
-                                    trust_remote_code=True,
-                                    token=os.environ.get("HF_TOKEN"),
-                                    low_cpu_mem_usage=True
-                    )
-                    print("Model loaded successfully!")
+DEFAULT_SYSTEM_PROMPT = "You are Kimi, an AI assistant created by Moonshot AI. You are helpful, harmless, and honest."
+
+# Initialize Inference Client
+client = None
+
+def init_client():
+        """Initialize the Hugging Face Inference Client"""""
+        global client
+        hf_token = os.environ.get("HF_TOKEN")
+        if hf_token:
+                    client = InferenceClient(token=hf_token)
+                    print("Inference client initialized successfully")
                     return True
+else:
+        print("Warning: HF_TOKEN not found. Please set it in Space secrets.")
+            return False
 
-        except Exception as e:
-                    print(f"Error loading model: {str(e)}")
-                    import traceback
-                    traceback.print_exc()
-                    return False
-
-            def generate_response(message, history, system_prompt, max_tokens, temperature, top_p):
-                    """Generate response from the model"""""
-                    global model, tokenizer
-
-                    if model is None or tokenizer is None:
-                                return "Model not loaded yet. Please wait..."
-
-                            try:
-                                        messages = [{"role": "system", "content": system_prompt or DEFAULT_SYSTEM_PROMPT}]
-
-                                        for h in history:
-                                                        if h[0]:
-                                                                            messages.append({"role": "user", "content": h[0]})
-                                                                        if h[1]:
-                                                                            messages.append({"role": "assistant", "content": h[1]})
-
-                                                                    messages.append({"role": "user", "content": message})
-
-                                                    inputs = tokenizer.apply_chat_template(
-                                                                    messages,
-                                                                    add_generation_prompt=True,
-                                                                    return_tensors="pt"
-                                                    ).to(model.device)
-
-        with torch.no_grad():
-                        outputs = model.generate(
-                                            inputs,
-                                            max_new_tokens=int(max_tokens),
-                                            temperature=float(temperature),
-                                            top_p=float(top_p),
-                                            do_sample=True,
-                                            pad_token_id=tokenizer.eos_token_id
-                        )
+def generate_response(message, history, system_prompt, max_tokens, temperature):
+        """Generate response using Hugging Face Inference API"""""
+    global client
 
-                    response = tokenizer.decode(
-                                    outputs[0][inputs.shape[1]:],
-                                    skip_special_tokens=True
-                    )
-        return response
+    if client is None:
+                if not init_client():
+                                return "Error: HF_TOKEN not configured. Please add it in Space settings."
 
-except Exception as e:
-        return f"Error: {str(e)}"
+            try:
+                        # Build messages
+                        messages = [{"role": "system", "content": system_prompt or DEFAULT_SYSTEM_PROMPT}]
 
-# Create interface
-print(f"\n===== Application Startup =====\n")
+                        for h in history:
+                                        if h[0]:
+                                                            messages.append({"role": "user", "content": h[0]})
+                                                        if h[1]:
+                                                            messages.append({"role": "assistant", "content": h[1]})
 
-model_loaded = load_model()
+                                    messages.append({"role": "user", "content": message})
 
-with gr.Blocks(title="Kimi-K2-Instruct") as iface:
-        gr.Markdown("# Kimi-K2-Instruct Chat")
+        # Call Inference API
+        response = client.chat_completion(
+                        model=MODEL_NAME,
+                        messages=messages,
+                        max_tokens=int(max_tokens),
+                        temperature=float(temperature)
+        )
 
-        if not model_loaded:
-                    gr.Markdown("**Warning:** Model failed to load.")
+        return response.choices[0].message.content
 
-                chatbot = gr.Chatbot(height=450)
+except Exception as e:
+        return f"Error: {str(e)}"
 
-    with gr.Row():
-                msg = gr.Textbox(placeholder="Type here...", label="Message", scale=4)
-                submit_btn = gr.Button("Send", variant="primary", scale=1)
+# Create interface
+print("===== Kimi K2 Thinking Dev =====")
+print(f"Using Inference API with model: {MODEL_NAME}")
+
+# Initialize client at startup
+client_ready = init_client()
 
-            with gr.Accordion("Settings", open=False):
-                        system_prompt = gr.Textbox(value=DEFAULT_SYSTEM_PROMPT, label="System Prompt")
-                        with gr.Row():
-                                        max_tokens = gr.Slider(64, 2048, 512, step=64, label="Max Tokens")
-                                        temperature = gr.Slider(0.1, 2.0, 0.6, step=0.1, label="Temperature")
-                                        top_p = gr.Slider(0.1, 1.0, 0.9, step=0.05, label="Top P")
+with gr.Blocks(title="Kimi-K2 Chat", theme=gr.themes.Soft()) as iface:
+        gr.Markdown("""
+            # 🤖 Kimi-K2 Instruct Chat
+                **Powered by Hugging Face Inference API**
+                    
+                        This space uses the Kimi-K2-Instruct model via API for efficient inference.
+                            """"")
 
-                                clear_btn = gr.Button("Clear")
+    if not client_ready:
+                gr.Markdown("⚠️ **Warning:** HF_TOKEN not found. Please configure it in Space secrets.")
 
-    def respond(message, history, system_prompt, max_tokens, temperature, top_p):
-                response = generate_response(message, history, system_prompt, max_tokens, temperature, top_p)
-                history.append((message, response))
-                return "", history
+    chatbot = gr.Chatbot(height=450, label="Chat")
 
-            msg.submit(respond, [msg, chatbot, system_prompt, max_tokens, temperature, top_p], [msg, chatbot])
-    submit_btn.click(respond, [msg, chatbot, system_prompt, max_tokens, temperature, top_p], [msg, chatbot])
+    with gr.Row():
+                msg = gr.Textbox(
+                                placeholder="Type your message here...",
+                                label="Your Message",
+                                scale=4,
+                                lines=2
+                )
+        submit_btn = gr.Button("Send 🚀", variant="primary", scale=1)
+
+    with gr.Accordion("⚙️ Settings", open=False):
+                system_prompt = gr.Textbox(
+                                value=DEFAULT_SYSTEM_PROMPT,
+                                label="System Prompt",
+                                lines=2
+                )
+        with gr.Row():
+                        max_tokens = gr.Slider(
+                            minimum=64,
+                                            maximum=2048,
+                            value=512,
+                            step=64,
+                            label="Max Tokens"
+        )
+            temperature = gr.Slider(
+                                minimum=0.1,
+                                maximum=2.0,
+                                value=0.7,
+                                step=0.1,
+                                label="Temperature"
+            )
+
+    clear_btn = gr.Button("🗑️ Clear Chat")
+
+    def respond(message, history, system_prompt, max_tokens, temperature):
+                if not message.strip():
+                                return "", history
+                            response = generate_response(message, history, system_prompt, max_tokens, temperature)
+        history.append((message, response))
+        return "", history
+
+    msg.submit(respond, [msg, chatbot, system_prompt, max_tokens, temperature], [msg, chatbot])
+    submit_btn.click(respond, [msg, chatbot, system_prompt, max_tokens, temperature], [msg, chatbot])
     clear_btn.click(lambda: [], None, chatbot)
 
 if __name__ == "__main__":
         iface.launch(server_name="0.0.0.0", server_port=7860)
-                    )
+            )
                         )
-                                                    )
-                    )
-                    )
+                )
+                )
+        )