Update app.py

app.py

@@ -3,7 +3,7 @@ import os
 import threading
 import time
 from pathlib import Path
-from huggingface_hub import hf_hub_download,login
+from huggingface_hub import hf_hub_download, login, list_repo_files
 
 # Try to import llama-cpp-python, fallback to instructions if not available
 try:
@@ -23,9 +23,9 @@ HF_FILENAME = "mmed-llama-alpaca-q4_k_m.gguf"
 
 hf_token = os.environ.get("HF_TOKEN")
 
-login(token = hf_token)
+if hf_token:
+    login(token=hf_token)
 
-login(token=hf_token)
 def find_gguf_file(directory="."):
     """Find GGUF files in the specified directory"""
     gguf_files = []
@@ -35,6 +35,19 @@ def find_gguf_file(directory="."):
                 gguf_files.append(os.path.join(root, file))
     return gguf_files
 
+def get_repo_gguf_files(repo_id=HF_REPO_ID):
+    """Get all GGUF files from the HuggingFace repository"""
+    try:
+        print(f"Fetching file list from {repo_id}...")
+        files = list_repo_files(repo_id=repo_id, token=hf_token)
+        gguf_files = [f for f in files if f.endswith('.gguf')]
+        print(f"Found {len(gguf_files)} GGUF files in repository")
+        return gguf_files, None
+    except Exception as e:
+        error_msg = f"Error fetching repository files: {str(e)}"
+        print(error_msg)
+        return [], error_msg
+
 def download_model_from_hf(repo_id=HF_REPO_ID, filename=HF_FILENAME):
     """Download GGUF model from HuggingFace Hub"""
     try:
@@ -76,7 +89,7 @@ def get_optimal_settings():
     
     return n_threads, n_gpu_layers
 
-def load_model_from_gguf(gguf_path=None, n_ctx=2048, use_hf_download=True):
+def load_model_from_gguf(gguf_path=None, filename=None, n_ctx=2048, use_hf_download=True):
     """Load the model from a GGUF file with automatic optimization"""
     global model, model_loaded
     
@@ -87,8 +100,10 @@ def load_model_from_gguf(gguf_path=None, n_ctx=2048, use_hf_download=True):
         # If no path provided, try different approaches
         if gguf_path is None:
             if use_hf_download:
+                # Use the specified filename or default
+                selected_filename = filename if filename else HF_FILENAME
                 # Try to download from HuggingFace first
-                gguf_path, error = download_model_from_hf()
+                gguf_path, error = download_model_from_hf(filename=selected_filename)
                 if error:
                     return False, f"❌ Failed to download from HuggingFace: {error}"
             else:
@@ -123,8 +138,9 @@ def load_model_from_gguf(gguf_path=None, n_ctx=2048, use_hf_download=True):
         )
         
         model_loaded = True
+        selected_filename = filename if filename else os.path.basename(gguf_path)
         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}\n📦 Source: {HF_REPO_ID}"
+        return True, f"✅ Model loaded successfully: {selected_filename}\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
@@ -195,11 +211,25 @@ def clear_chat():
     """Clear the chat history"""
     return [], ""
 
-def load_model_interface(context_size, use_hf_download):
+def load_model_interface(context_size, selected_model):
     """Interface function to load model with configurable context size"""
-    success, message = load_model_from_gguf(gguf_path=None, n_ctx=int(context_size), use_hf_download=use_hf_download)
+    success, message = load_model_from_gguf(gguf_path=None, filename=selected_model, n_ctx=int(context_size), use_hf_download=True)
     return message
 
+def refresh_model_list():
+    """Refresh the list of available GGUF models from the repository"""
+    gguf_files, error = get_repo_gguf_files()
+    if error:
+        return gr.Dropdown(choices=["Error loading models"], value="Error loading models")
+    
+    if not gguf_files:
+        return gr.Dropdown(choices=["No GGUF files found"], value="No GGUF files found")
+    
+    # Set default value to the original default file if it exists
+    default_value = HF_FILENAME if HF_FILENAME in gguf_files else gguf_files[0]
+    
+    return gr.Dropdown(choices=gguf_files, value=default_value)
+
 def get_available_gguf_files():
     """Get list of available GGUF files"""
     gguf_files = find_gguf_file()
@@ -220,6 +250,15 @@ def create_interface():
     # Check for available models
     availability_status = check_model_availability()
     
+    # Get initial list of GGUF files from repository
+    gguf_files, error = get_repo_gguf_files()
+    if error or not gguf_files:
+        initial_choices = ["Error loading models" if error else "No GGUF files found"]
+        initial_value = initial_choices[0]
+    else:
+        initial_choices = gguf_files
+        initial_value = HF_FILENAME if HF_FILENAME in gguf_files else gguf_files[0]
+    
     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;">
@@ -256,15 +295,16 @@ def create_interface():
                 # Model loading section
                 gr.HTML("<h3>🔧 Model Control</h3>")
                 
-                # Model source selection
-                use_hf_download = gr.Checkbox(
-                    value=True,
-                    label="Download from HuggingFace",
-                    info="Uncheck to use local GGUF files"
-                )
+                # 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>Repository:</strong> {HF_REPO_ID}</p>")
-                gr.HTML(f"<p style='font-size: 0.9em; color: #666;'><strong>File:</strong> {HF_FILENAME}</p>")
+                # Model selection dropdown
+                model_dropdown = gr.Dropdown(
+                    choices=initial_choices,
+                    value=initial_value,
+                    label="Select GGUF Model",
+                    info="Choose from available models in the repository",
+                    interactive=True
+                )
                 
                 load_btn = gr.Button("Load Model", variant="primary", size="lg")
                 model_status = gr.Textbox(
@@ -346,7 +386,7 @@ def create_interface():
         # Event handlers
         load_btn.click(
             load_model_interface,
-            inputs=[context_size, use_hf_download],
+            inputs=[context_size, model_dropdown],
             outputs=model_status
         )