Improve error handling and use Q2_K quantization for faster loading

app.py

@@ -5,32 +5,66 @@ import gradio as gr
 from llama_cpp import Llama
 import logging
 from huggingface_hub import hf_hub_download
+import os
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 MODEL_REPO = "Jiunsong/supergemma4-26b-uncensored-gguf-v2"
-MODEL_FILE = "supergemma4-26b-uncensored-Q4_K_M.gguf"
+# Try Q2_K for smaller size and faster loading
+MODEL_FILE = "supergemma4-26b-uncensored-Q2_K.gguf"
 
 logger.info(f"Loading model: {MODEL_REPO}/{MODEL_FILE}")
+logger.info(f"This may take 5-10 minutes for first load...")
+
+llm = None
 
 try:
-    model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE, repo_type="model")
+    logger.info("Downloading model from HuggingFace...")
+    model_path = hf_hub_download(
+        repo_id=MODEL_REPO, 
+        filename=MODEL_FILE, 
+        repo_type="model",
+        resume_download=True
+    )
     logger.info(f"Model downloaded to: {model_path}")
-    llm = Llama(model_path=model_path, n_ctx=4096, n_threads=8, n_gpu_layers=0, verbose=False)
-    logger.info("Model loaded successfully on CPU")
+    logger.info(f"Model file size: {os.path.getsize(model_path) / (1024**3):.2f} GB")
+    
+    logger.info("Loading model into memory...")
+    llm = Llama(
+        model_path=model_path,
+        n_ctx=2048,  # Reduced context for faster inference
+        n_threads=4,  # Reduced threads
+        n_gpu_layers=0,  # CPU only
+        verbose=True,
+        n_batch=512
+    )
+    logger.info("✅ Model loaded successfully on CPU!")
+    
 except Exception as e:
-    logger.error(f"Error loading model: {str(e)}")
+    logger.error(f"❌ Error loading model: {str(e)}")
+    logger.error(f"Full error: {repr(e)}")
     llm = None
 
 def generate_text(prompt, max_tokens=500, temperature=0.7, top_p=0.9, top_k=40):
     if llm is None:
-        return "Error: Model not loaded"
+        return "❌ Error: Model not loaded. Check Space logs for details."
     try:
-        response = llm(prompt, max_tokens=int(max_tokens), temperature=float(temperature), 
-                      top_p=float(top_p), top_k=int(top_k), stop=["</s>"], echo=False)
-        return response['choices'][0]['text'].strip()
+        logger.info(f"Generating: {prompt[:50]}...")
+        response = llm(
+            prompt, 
+            max_tokens=int(max_tokens), 
+            temperature=float(temperature), 
+            top_p=float(top_p), 
+            top_k=int(top_k), 
+            stop=["</s>", "\n\n\n"], 
+            echo=False
+        )
+        result = response['choices'][0]['text'].strip()
+        logger.info(f"Generated {len(result)} characters")
+        return result
     except Exception as e:
+        logger.error(f"Generation error: {str(e)}")
         return f"Error: {str(e)}"
 
 def generate_code(prompt, max_tokens=500, temperature=0.2, top_p=0.95):
@@ -39,24 +73,52 @@ def generate_code(prompt, max_tokens=500, temperature=0.2, top_p=0.95):
 
 def chat(message, history, max_tokens=500, temperature=0.7):
     if llm is None:
-        return "Error: Model not loaded"
+        return "❌ Error: Model not loaded"
     conversation = ""
     for user_msg, assistant_msg in history:
-        conversation += f"### User:\n{user_msg}\n\n### Assistant:\n{assistant_msg}\n\n"
-    conversation += f"### User:\n{message}\n\n### Assistant:\n"
-    response = llm(conversation, max_tokens=int(max_tokens), temperature=float(temperature), 
-                  top_p=0.9, top_k=40, stop=["### User:"], echo=False)
+        conversation += f"User: {user_msg}\nAssistant: {assistant_msg}\n\n"
+    conversation += f"User: {message}\nAssistant: "
+    response = llm(
+        conversation, 
+        max_tokens=int(max_tokens), 
+        temperature=float(temperature), 
+        top_p=0.9, 
+        top_k=40, 
+        stop=["User:", "</s>"], 
+        echo=False
+    )
     return response['choices'][0]['text'].strip()
 
+# Create UI
 with gr.Blocks(title="SuperGemma4-26B Uncensored", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🚀 SuperGemma4-26B Uncensored (CPU)\n\n26B parameter uncensored model running on CPU")
+    gr.Markdown(f"""
+    # 🚀 SuperGemma4-26B Uncensored (CPU)
+    
+    **Status**: {'✅ Model Loaded' if llm else '❌ Model Loading Failed'}
+    
+    26B parameter uncensored model running on CPU with Q2_K quantization
+    
+    ⚠️ **Note**: First load takes 5-10 minutes. Please be patient!
+    """)
+    
+    if llm is None:
+        gr.Markdown("""
+        ### ⚠️ Model Loading Error
+        
+        The model failed to load. Possible reasons:
+        1. Model file is still downloading (check Space logs)
+        2. Insufficient memory
+        3. Model file not found
+        
+        **Check the Logs tab** in your Space for detailed error messages.
+        """)
     
     with gr.Tabs():
         with gr.Tab("💬 Chat"):
             chatbot = gr.Chatbot(height=400)
-            msg = gr.Textbox(label="Message")
+            msg = gr.Textbox(label="Message", placeholder="Ask anything...")
             with gr.Row():
-                chat_max_tokens = gr.Slider(100, 1000, value=500, label="Max Tokens")
+                chat_max_tokens = gr.Slider(100, 1000, value=300, label="Max Tokens")
                 chat_temperature = gr.Slider(0.1, 2.0, value=0.7, label="Temperature")
             with gr.Row():
                 submit = gr.Button("Send", variant="primary")
@@ -73,8 +135,8 @@ with gr.Blocks(title="SuperGemma4-26B Uncensored", theme=gr.themes.Soft()) as de
         with gr.Tab("💻 Generate Code"):
             with gr.Row():
                 with gr.Column():
-                    gen_prompt = gr.Textbox(label="Prompt", lines=5)
-                    gen_max_tokens = gr.Slider(100, 1000, value=500, label="Max Tokens")
+                    gen_prompt = gr.Textbox(label="Prompt", lines=5, placeholder="Write a Python function to...")
+                    gen_max_tokens = gr.Slider(100, 1000, value=400, label="Max Tokens")
                     gen_temperature = gr.Slider(0.1, 1.0, value=0.2, label="Temperature")
                     gen_top_p = gr.Slider(0.1, 1.0, value=0.95, label="Top P")
                     gen_button = gr.Button("Generate", variant="primary")
@@ -85,8 +147,8 @@ with gr.Blocks(title="SuperGemma4-26B Uncensored", theme=gr.themes.Soft()) as de
         with gr.Tab("📝 Generate Text"):
             with gr.Row():
                 with gr.Column():
-                    text_prompt = gr.Textbox(label="Prompt", lines=5)
-                    text_max_tokens = gr.Slider(100, 1000, value=500, label="Max Tokens")
+                    text_prompt = gr.Textbox(label="Prompt", lines=5, placeholder="Write about...")
+                    text_max_tokens = gr.Slider(100, 1000, value=400, label="Max Tokens")
                     text_temperature = gr.Slider(0.1, 2.0, value=0.7, label="Temperature")
                     text_top_p = gr.Slider(0.1, 1.0, value=0.9, label="Top P")
                     text_top_k = gr.Slider(1, 100, value=40, label="Top K")
@@ -94,6 +156,13 @@ with gr.Blocks(title="SuperGemma4-26B Uncensored", theme=gr.themes.Soft()) as de
                 with gr.Column():
                     text_output = gr.Textbox(label="Generated Text", lines=20)
             text_button.click(generate_text, [text_prompt, text_max_tokens, text_temperature, text_top_p, text_top_k], text_output)
+    
+    gr.Markdown("""
+    ---
+    **Model**: SuperGemma4-26B-Uncensored (Q2_K) | **Hardware**: CPU | **Powered by**: llama.cpp
+    
+    ⚠️ CPU inference is slower (~1-3 tokens/second). Be patient with responses.
+    """)
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)