Update app.py

app.py

@@ -3,8 +3,7 @@ import os
 import threading
 import time
 from pathlib import Path
-from huggingface_hub import login
-
+from huggingface_hub import hf_hub_download
 
 # Try to import llama-cpp-python, fallback to instructions if not available
 try:
@@ -14,15 +13,14 @@ except ImportError:
     LLAMA_CPP_AVAILABLE = False
     print("llama-cpp-python not installed. Please install it with: pip install llama-cpp-python")
 
-hf_token = os.environ.get("HF_TOKEN")
-
-login(token = hf_token)
-
-
 # Global variables for model
 model = None
 model_loaded = False
 
+# HuggingFace repository information
+HF_REPO_ID = "Axcel1/MMed-llama-alpaca-Q4_K_M-GGUF"
+HF_FILENAME = "mmed-llama-alpaca-q4_k_m.gguf"
+
 def find_gguf_file(directory="."):
     """Find GGUF files in the specified directory"""
     gguf_files = []
@@ -32,11 +30,32 @@ def find_gguf_file(directory="."):
                 gguf_files.append(os.path.join(root, file))
     return gguf_files
 
+def download_model_from_hf(repo_id=HF_REPO_ID, filename=HF_FILENAME):
+    """Download GGUF model from HuggingFace Hub"""
+    try:
+        print(f"Downloading model from {repo_id}/{filename}...")
+        gguf_path = hf_hub_download(
+            repo_id=repo_id,
+            filename=filename,
+            cache_dir="./models",
+            resume_download=True  # Resume partial downloads
+        )
+        print(f"Model downloaded to: {gguf_path}")
+        return gguf_path, None
+    except Exception as e:
+        error_msg = f"Error downloading model: {str(e)}"
+        print(error_msg)
+        return None, error_msg
+
 def get_optimal_settings():
     """Get optimal CPU threads and GPU layers automatically"""
     # Auto-detect CPU threads (use all available cores)
     n_threads = os.cpu_count()
     
+    # For Hugging Face Spaces, limit threads to avoid resource issues
+    if n_threads and n_threads > 4:
+        n_threads = 4
+    
     # Auto-detect GPU layers (try to use GPU if available)
     n_gpu_layers = 0
     try:
@@ -52,59 +71,28 @@ def get_optimal_settings():
     
     return n_threads, n_gpu_layers
 
-def load_model_from_huggingface(repo_id, filename, n_ctx=2048):
-    """Load the model from Hugging Face repository"""
-    global model, model_loaded
-    
-    if not LLAMA_CPP_AVAILABLE:
-        return False, "llama-cpp-python not installed. Please install it with: pip install llama-cpp-python"
-    
-    try:
-        print(f"Loading model from Hugging Face: {repo_id}/{filename}")
-        
-        # Get optimal settings automatically
-        n_threads, n_gpu_layers = get_optimal_settings()
-        print(f"Auto-detected settings: {n_threads} CPU threads, {n_gpu_layers} GPU layers")
-        
-        # Load model from Hugging Face with optimized settings
-        model = Llama.from_pretrained(
-            repo_id=repo_id,
-            filename=filename,
-            n_ctx=n_ctx,  # Context window (configurable)
-            n_threads=n_threads,  # CPU threads (auto-detected)
-            n_gpu_layers=n_gpu_layers,  # Number of layers to offload to GPU (auto-detected)
-            verbose=False,
-            chat_format="chatml",  # Use Llama-3 chat format
-            n_batch=512,  # Batch size for prompt processing
-            use_mlock=True,  # Keep model in memory
-            use_mmap=True,  # Use memory mapping
-        )
-        
-        model_loaded = True
-        print("Model loaded successfully!")
-        return True, f"✅ Model loaded successfully from {repo_id}/{filename}\n📊 Context: {n_ctx} tokens\n🖥️ CPU Threads: {n_threads}\n🎮 GPU Layers: {n_gpu_layers}"
-        
-    except Exception as e:
-        model_loaded = False
-        error_msg = f"Error loading model: {str(e)}"
-        print(error_msg)
-        return False, f"❌ {error_msg}"
-
-def load_model_from_gguf(gguf_path=None, n_ctx=2048):
-    """Load the model from a local GGUF file with automatic optimization"""
+def load_model_from_gguf(gguf_path=None, n_ctx=2048, use_hf_download=True):
+    """Load the model from a GGUF file with automatic optimization"""
     global model, model_loaded
     
     if not LLAMA_CPP_AVAILABLE:
         return False, "llama-cpp-python not installed. Please install it with: pip install llama-cpp-python"
     
     try:
-        # If no path provided, try to find GGUF files
+        # If no path provided, try different approaches
         if gguf_path is None:
-            gguf_files = find_gguf_file()
-            if not gguf_files:
-                return False, "No GGUF files found in the repository"
-            gguf_path = gguf_files[0]  # Use the first one found
-            print(f"Found GGUF file: {gguf_path}")
+            if use_hf_download:
+                # Try to download from HuggingFace first
+                gguf_path, error = download_model_from_hf()
+                if error:
+                    return False, f"❌ Failed to download from HuggingFace: {error}"
+            else:
+                # Try to find local GGUF files
+                gguf_files = find_gguf_file()
+                if not gguf_files:
+                    return False, "No GGUF files found in the repository"
+                gguf_path = gguf_files[0]  # Use the first one found
+                print(f"Found local GGUF file: {gguf_path}")
         
         # Check if file exists
         if not os.path.exists(gguf_path):
@@ -116,22 +104,22 @@ def load_model_from_gguf(gguf_path=None, n_ctx=2048):
         n_threads, n_gpu_layers = get_optimal_settings()
         print(f"Auto-detected settings: {n_threads} CPU threads, {n_gpu_layers} GPU layers")
         
-        # Load model with optimized settings
+        # Load model with optimized settings for Hugging Face Spaces
         model = Llama(
             model_path=gguf_path,
             n_ctx=n_ctx,  # Context window (configurable)
-            n_threads=n_threads,  # CPU threads (auto-detected)
-            n_gpu_layers=n_gpu_layers,  # Number of layers to offload to GPU (auto-detected)
+            n_threads=n_threads,  # CPU threads (limited for Spaces)
+            n_gpu_layers=n_gpu_layers,  # Number of layers to offload to GPU
             verbose=False,
             chat_format="llama-3",  # Use Llama-3 chat format
-            n_batch=512,  # Batch size for prompt processing
-            use_mlock=True,  # Keep model in memory
+            n_batch=256,  # Smaller batch size for Spaces
+            use_mlock=False,  # Disabled for Spaces compatibility
             use_mmap=True,  # Use memory mapping
         )
         
         model_loaded = True
         print("Model loaded successfully!")
-        return True, f"✅ Model loaded successfully from {os.path.basename(gguf_path)}\n📊 Context: {n_ctx} tokens\n🖥️ CPU Threads: {n_threads}\n🎮 GPU Layers: {n_gpu_layers}"
+        return True, f"✅ Model loaded successfully from {os.path.basename(gguf_path)}\n📊 Context: {n_ctx} tokens\n🖥️ CPU Threads: {n_threads}\n🎮 GPU Layers: {n_gpu_layers}\n📦 Source: {HF_REPO_ID}"
         
     except Exception as e:
         model_loaded = False
@@ -202,34 +190,39 @@ def clear_chat():
     """Clear the chat history"""
     return [], ""
 
-def load_model_interface(source_type, gguf_path, repo_id, filename, context_size):
+def load_model_interface(context_size, use_hf_download):
     """Interface function to load model with configurable context size"""
-    if source_type == "Hugging Face":
-        success, message = load_model_from_huggingface(repo_id, filename, n_ctx=int(context_size))
-    else:  # Local file
-        success, message = load_model_from_gguf(gguf_path, n_ctx=int(context_size))
+    success, message = load_model_from_gguf(gguf_path=None, n_ctx=int(context_size), use_hf_download=use_hf_download)
     return message
 
 def get_available_gguf_files():
     """Get list of available GGUF files"""
     gguf_files = find_gguf_file()
     if not gguf_files:
-        return ["No GGUF files found"]
+        return ["No local GGUF files found"]
     return [os.path.basename(f) for f in gguf_files]
 
+def check_model_availability():
+    """Check if model is available locally or needs to be downloaded"""
+    local_files = find_gguf_file()
+    if local_files:
+        return f"Local GGUF files found: {len(local_files)}"
+    else:
+        return "No local GGUF files found. Will download from HuggingFace."
+
 # Create the Gradio interface
 def create_interface():
-    # Get available GGUF files
-    gguf_files = find_gguf_file()
-    gguf_choices = [os.path.basename(f) for f in gguf_files] if gguf_files else ["No GGUF files found"]
+    # Check for available models
+    availability_status = check_model_availability()
     
-    with gr.Blocks(title="Llama-3-8B GGUF Chatbot", theme=gr.themes.Soft()) as demo:
+    with gr.Blocks(title="MMed-Llama-Alpaca GGUF Chatbot", theme=gr.themes.Soft()) as demo:
         gr.HTML("""
             <h1 style="text-align: center; color: #2E86AB; margin-bottom: 30px;">
                 🦙 MMed-Llama-Alpaca GGUF Chatbot
             </h1>
             <p style="text-align: center; color: #666; margin-bottom: 30px;">
-                Chat with the MMed-Llama-Alpaca model (Q4_K_M quantized) for medical assistance!
+                Chat with the MMed-Llama-Alpaca model (Q4_K_M quantized) for medical assistance!<br>
+                <strong>⚠️ This is for educational purposes only. Always consult healthcare professionals for medical advice.</strong>
             </p>
         """)
         
@@ -246,7 +239,7 @@ def create_interface():
                 
                 with gr.Row():
                     msg = gr.Textbox(
-                        placeholder="Type your message here...",
+                        placeholder="Type your medical question here...",
                         container=False,
                         scale=7,
                         show_label=False
@@ -259,59 +252,30 @@ def create_interface():
                 gr.HTML("<h3>🔧 Model Control</h3>")
                 
                 # Model source selection
-                source_type = gr.Radio(
-                    choices=["Hugging Face", "Local File"],
-                    value="Hugging Face",
-                    label="Model Source",
-                    info="Choose where to load the model from"
+                use_hf_download = gr.Checkbox(
+                    value=True,
+                    label="Download from HuggingFace",
+                    info="Uncheck to use local GGUF files"
                 )
                 
-                # Hugging Face settings
-                with gr.Group(visible=True) as hf_group:
-                    gr.HTML("<h4>🤗 Hugging Face Settings</h4>")
-                    repo_id = gr.Textbox(
-                        value="Axcel1/MMed-llama-alpaca-Q4_K_M-GGUF",
-                        label="Repository ID",
-                        info="e.g., username/repo-name"
-                    )
-                    filename = gr.Textbox(
-                        value="mmed-llama-alpaca-q4_k_m.gguf",
-                        label="Filename",
-                        info="GGUF filename in the repository"
-                    )
-                
-                # Local file settings
-                with gr.Group(visible=False) as local_group:
-                    gr.HTML("<h4>📁 Local File Settings</h4>")
-                    if gguf_files:
-                        gguf_dropdown = gr.Dropdown(
-                            choices=gguf_choices,
-                            value=gguf_choices[0] if gguf_choices[0] != "No GGUF files found" else None,
-                            label="Select GGUF File",
-                            info="Choose which GGUF file to load"
-                        )
-                    else:
-                        gguf_dropdown = gr.Textbox(
-                            value="No GGUF files found in repository",
-                            label="GGUF File",
-                            interactive=False
-                        )
+                gr.HTML(f"<p style='font-size: 0.9em; color: #666;'><strong>Repository:</strong> {HF_REPO_ID}</p>")
+                gr.HTML(f"<p style='font-size: 0.9em; color: #666;'><strong>File:</strong> {HF_FILENAME}</p>")
                 
                 load_btn = gr.Button("Load Model", variant="primary", size="lg")
                 model_status = gr.Textbox(
                     label="Status",
-                    value="Model not loaded. Configure settings and click 'Load Model'.\n⚙️ Auto-optimized: CPU threads & GPU layers auto-detected\n📝 Context size can be configured in Generation Settings",
+                    value=f"Model not loaded.\n{availability_status}\n⚙️ Auto-optimized: CPU threads & GPU layers auto-detected\n📝 Context size can be configured below",
                     interactive=False,
-                    max_lines=5
+                    max_lines=6
                 )
                 
                 # Generation parameters
                 gr.HTML("<h3>⚙️ Generation Settings</h3>")
                 
-                # Context size (now as a slider)
+                # Context size (limited for Spaces)
                 context_size = gr.Slider(
                     minimum=512,
-                    maximum=8192,
+                    maximum=4096,
                     value=2048,
                     step=256,
                     label="Context Size",
@@ -320,7 +284,7 @@ def create_interface():
                 
                 max_tokens = gr.Slider(
                     minimum=50,
-                    maximum=2048,
+                    maximum=1024,
                     value=512,
                     step=50,
                     label="Max Tokens",
@@ -354,12 +318,13 @@ def create_interface():
                 # Information section
                 gr.HTML("""
                     <h3>ℹ️ About</h3>
+                    <p><strong>Model:</strong> MMed-Llama-Alpaca</p>
+                    <p><strong>Quantization:</strong> Q4_K_M</p>
                     <p><strong>Format:</strong> GGUF (optimized)</p>
                     <p><strong>Backend:</strong> llama-cpp-python</p>
                     <p><strong>Features:</strong> CPU/GPU support, streaming</p>
-                    <p><strong>Memory:</strong> Optimized usage</p>
+                    <p><strong>Specialty:</strong> Medical assistance</p>
                     <p><strong>Auto-Optimization:</strong> CPU threads & GPU layers detected automatically</p>
-                    <p><strong>Sources:</strong> Hugging Face Hub or Local Files</p>
                 """)
                 
                 if not LLAMA_CPP_AVAILABLE:
@@ -374,21 +339,9 @@ def create_interface():
                     """)
         
         # Event handlers
-        def toggle_source_visibility(source_type):
-            if source_type == "Hugging Face":
-                return gr.update(visible=True), gr.update(visible=False)
-            else:
-                return gr.update(visible=False), gr.update(visible=True)
-        
-        source_type.change(
-            toggle_source_visibility,
-            inputs=source_type,
-            outputs=[hf_group, local_group]
-        )
-        
         load_btn.click(
             load_model_interface,
-            inputs=[source_type, gguf_dropdown, repo_id, filename, context_size],
+            inputs=[context_size, use_hf_download],
             outputs=model_status
         )
         
@@ -415,7 +368,7 @@ if __name__ == "__main__":
     # Create and launch the interface
     demo = create_interface()
     
-    # Launch with appropriate settings for Hugging Face Spaces
+    # Launch with settings optimized for Hugging Face Spaces
     demo.launch(
         server_name="0.0.0.0",
         server_port=7860,