瀏覽代碼

adding functionality to delete the models if there is part of the model downloaded

cadenmackenzie 8 月之前
父節點
當前提交
619df1d49e
共有 4 個文件被更改,包括 138 次插入1 次删除
  1. 55 0
      exo/api/chatgpt_api.py
  2. 26 0
      exo/tinychat/index.css
  3. 9 1
      exo/tinychat/index.html
  4. 48 0
      exo/tinychat/index.js

+ 55 - 0
exo/api/chatgpt_api.py

@@ -19,6 +19,8 @@ from exo.orchestration import Node
 from exo.models import build_base_shard, model_cards, get_repo, pretty_name
 from typing import Callable, Optional
 from exo.download.hf.hf_shard_download import HFShardDownloader
+import shutil
+from exo.download.hf.hf_helpers import get_hf_home, get_repo_root
 
 class Message:
   def __init__(self, role: str, content: Union[str, List[Dict[str, Union[str, Dict[str, str]]]]]):
@@ -176,6 +178,7 @@ class ChatGPTAPI:
     cors.add(self.app.router.add_get("/modelpool", self.handle_model_support), {"*": cors_options})
     cors.add(self.app.router.add_get("/healthcheck", self.handle_healthcheck), {"*": cors_options})
     cors.add(self.app.router.add_post("/quit", self.handle_quit), {"*": cors_options})
+    cors.add(self.app.router.add_delete("/models/{model_name}", self.handle_delete_model), {"*": cors_options})
 
     if "__compiled__" not in globals():
       self.static_dir = Path(__file__).parent.parent/"tinychat"
@@ -415,6 +418,58 @@ class ChatGPTAPI:
       deregistered_callback = self.node.on_token.deregister(callback_id)
       if DEBUG >= 2: print(f"Deregister {callback_id=} {deregistered_callback=}")
 
+  async def handle_delete_model(self, request):
+    try:
+        model_name = request.match_info.get('model_name')
+        if DEBUG >= 2: print(f"Attempting to delete model: {model_name}")
+        
+        if not model_name or model_name not in model_cards:
+            return web.json_response(
+                {"detail": f"Invalid model name: {model_name}"}, 
+                status=400
+            )
+
+        shard = build_base_shard(model_name, self.inference_engine_classname)
+        if not shard:
+            return web.json_response(
+                {"detail": "Could not build shard for model"}, 
+                status=400
+            )
+
+        repo_id = get_repo(shard.model_id, self.inference_engine_classname)
+        if DEBUG >= 2: print(f"Repo ID for model: {repo_id}")
+        
+        # Get the HF cache directory using the helper function
+        hf_home = get_hf_home()
+        cache_dir = get_repo_root(repo_id)
+        
+        if DEBUG >= 2: print(f"Looking for model files in: {cache_dir}")
+        
+        if os.path.exists(cache_dir):
+            if DEBUG >= 2: print(f"Found model files at {cache_dir}, deleting...")
+            try:
+                shutil.rmtree(cache_dir)
+                return web.json_response({
+                    "status": "success", 
+                    "message": f"Model {model_name} deleted successfully",
+                    "path": str(cache_dir)
+                })
+            except Exception as e:
+                return web.json_response({
+                    "detail": f"Failed to delete model files: {str(e)}"
+                }, status=500)
+        else:
+            return web.json_response({
+                "detail": f"Model files not found at {cache_dir}"
+            }, status=404)
+            
+    except Exception as e:
+        print(f"Error in handle_delete_model: {str(e)}")
+        traceback.print_exc()
+        return web.json_response({
+            "detail": f"Server error: {str(e)}"
+        }, status=500)
+
   async def run(self, host: str = "0.0.0.0", port: int = 52415):
     runner = web.AppRunner(self.app)
     await runner.setup()

+ 26 - 0
exo/tinychat/index.css

@@ -557,4 +557,30 @@ main {
   font-size: 0.8em;
   color: var(--secondary-color-transparent);
   opacity: 0.8;
+}
+
+.model-header {
+    display: flex;
+    justify-content: space-between;
+    align-items: center;
+    margin-bottom: 4px;
+}
+
+.model-delete-button {
+    background: none;
+    border: none;
+    color: var(--red-color);
+    padding: 4px 8px;
+    cursor: pointer;
+    transition: all 0.2s ease;
+    opacity: 0.7;
+}
+
+.model-delete-button:hover {
+    opacity: 1;
+    transform: scale(1.1);
+}
+
+.model-option:hover .model-delete-button {
+    opacity: 1;
 }

+ 9 - 1
exo/tinychat/index.html

@@ -31,7 +31,15 @@
         <div class="model-option" 
              :class="{ 'selected': cstate.selectedModel === key }"
              @click="cstate.selectedModel = key">
-            <div class="model-name" x-text="model.name"></div>
+            <div class="model-header">
+                <div class="model-name" x-text="model.name"></div>
+                <button 
+                    @click.stop="deleteModel(key, model)"
+                    class="model-delete-button"
+                    x-show="model.download_percentage > 0">
+                    <i class="fas fa-trash"></i>
+                </button>
+            </div>
             <div class="model-info">
                 <div class="model-progress">
                     <template x-if="model.download_percentage != null">

+ 48 - 0
exo/tinychat/index.js

@@ -435,6 +435,54 @@ document.addEventListener("alpine:init", () => {
         }, 30 * 1000);
       }
     },
+
+    async deleteModel(modelName, model) {
+        const downloadedSize = model.total_downloaded || 0;
+        const sizeMessage = downloadedSize > 0 ? 
+            `This will free up ${this.formatBytes(downloadedSize)} of space.` :
+            'This will remove any partially downloaded files.';
+        
+        if (!confirm(`Are you sure you want to delete ${model.name}? ${sizeMessage}`)) {
+            return;
+        }
+
+        try {
+            const response = await fetch(`${window.location.origin}/models/${modelName}`, {
+                method: 'DELETE',
+                headers: {
+                    'Content-Type': 'application/json'
+                }
+            });
+
+            const data = await response.json();
+            
+            if (!response.ok) {
+                throw new Error(data.detail || 'Failed to delete model');
+            }
+
+            // Update the model status in the UI
+            if (this.models[modelName]) {
+                this.models[modelName].downloaded = false;
+                this.models[modelName].download_percentage = 0;
+                this.models[modelName].total_downloaded = 0;
+            }
+
+            // If this was the selected model, switch to a different one
+            if (this.cstate.selectedModel === modelName) {
+                const availableModel = Object.keys(this.models).find(key => this.models[key].downloaded);
+                this.cstate.selectedModel = availableModel || 'llama-3.2-1b';
+            }
+
+            // Show success message
+            console.log(`Model deleted successfully from: ${data.path}`);
+
+            // Refresh the model list
+            await this.populateSelector();
+        } catch (error) {
+            console.error('Error deleting model:', error);
+            this.setError(error.message || 'Failed to delete model');
+        }
+    }
   }));
 });