Skip to content

Commit

Permalink
Merge pull request #30 from NexaAI/david/bugfix
Browse files Browse the repository at this point in the history 
change to nexa run
  • Loading branch information
@zhiyuan8
zhiyuan8 authored Aug 23, 2024
2 parents 9aa5d44 + bfed39c commit c292fab
Show file tree
Hide file tree
Showing 25 changed files with 606 additions and 361 deletions.
261 changes: 115 additions & 146 deletions nexa/cli/entry.py
View file

Large diffs are not rendered by default.

92 changes: 89 additions & 3 deletions nexa/constants.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,12 @@
"phi3v": "Phi-3-vision-128k-instruct:onnx-cpu-int4",
"sd1-5": "stable-diffusion-v1-5:onnx-cpu-fp32",
"lcm-dreamshaper": "lcm-dreamshaper-v7:onnx-cpu-fp32",
"whisper": "whisper-tiny:onnx-cpu-fp32",
"whisper-tiny.en": "whisper-tiny.en:onnx-cpu-fp32",
"whisper-tiny": "whisper-tiny:onnx-cpu-fp32",
"whisper-small.en": "whisper-small.en:onnx-cpu-fp32",
"whisper-small": "whisper-small:onnx-cpu-fp32",
"whisper-base.en": "whisper-base.en:onnx-cpu-fp32",
"whisper-base": "whisper-base:onnx-cpu-fp32",
"ljspeech": "ljspeech-jets:onnx-cpu-fp32",
}

Expand All @@ -86,10 +91,14 @@
"whisper-large": "whisper-large:bin-large-v3",
"whisper-tiny": "whisper-tiny:bin-tiny",
"faster-whisper-tiny": "faster-whisper-tiny:bin-cpu-fp16",
"faster-whisper-tiny.en": "faster-whisper-tiny.en:bin-cpu-fp16",
"faster-whisper-small": "faster-whisper-small:bin-cpu-fp16",
"faster-whisper-small.en": "faster-whisper-small.en:bin-cpu-fp16",
"faster-whisper-medium": "faster-whisper-medium:bin-cpu-fp16",
"faster-whisper-medium.en": "faster-whisper-medium.en:bin-cpu-fp16",
"faster-whisper-base": "faster-whisper-base:bin-cpu-fp16",
"faster-whisper-large": "faster-whisper-large:bin-cpu-fp16",
"faster-whisper-base.en": "faster-whisper-base.en:bin-cpu-fp16",
"faster-whisper-large": "faster-whisper-large-v3:bin-cpu-fp16",
}

NEXA_RUN_MODEL_MAP_FUNCTION_CALLING = {
Expand Down Expand Up @@ -129,6 +138,8 @@
**NEXA_RUN_MODEL_MAP_TEXT,
**NEXA_RUN_MODEL_MAP_VLM,
**NEXA_RUN_MODEL_MAP_IMAGE,
**NEXA_RUN_MODEL_MAP_VOICE,
**NEXA_RUN_MODEL_MAP_FUNCTION_CALLING,
}

NEXA_RUN_CHAT_TEMPLATE_MAP = {
Expand Down Expand Up @@ -205,4 +216,79 @@
"task": "transcribe",
"temperature": 0.0,
"compute_type": "default"
}
}

NEXA_OFFICIAL_MODELS_TYPE = {
'gemma-2b': 'NLP',
'Llama-2-7b-chat': 'NLP',
'Llama-2-7b': 'NLP',
'Meta-Llama-3-8B-Instruct': 'NLP',
'Meta-Llama-3.1-8B-Instruct': 'NLP',
'Mistral-7B-Instruct-v0.3': 'NLP',
'Mistral-7B-Instruct-v0.2': 'NLP',
'Phi-3-mini-128k-instruct': 'NLP',
'Phi-3-mini-4k-instruct': 'NLP',
'CodeQwen1.5-7B-Instruct': 'NLP',
'Qwen2-0.5B-Instruct': 'NLP',
'Qwen2-1.5B-Instruct': 'NLP',
'Qwen2-7B-Instruct': 'NLP',
'codegemma-2b': 'NLP',
'gemma-1.1-2b-instruct': 'NLP',
'gemma-2b-instruct': 'NLP',
'gemma-2-9b-instruct': 'NLP',
'gemma-1.1-7b-instruct': 'NLP',
'gemma-7b-instruct': 'NLP',
'gemma-7b': 'NLP',
'Qwen2-1.5B': 'NLP',
'codegemma-7b': 'NLP',
'TinyLlama-1.1B-Chat-v1.0': 'NLP',
'CodeLlama-7b-Instruct': 'NLP',
'gpt2': 'NLP',
'CodeLlama-7b': 'NLP',
'CodeLlama-7b-Python': 'NLP',
'Qwen1.5-7B-Instruct': 'NLP',
'Qwen1.5-7B': 'NLP',
'Phi-2': 'NLP',
'deepseek-coder-1.3b-instruct': 'NLP',
'deepseek-coder-1.3b-base': 'NLP',
'deepseek-coder-6.7b-instruct': 'NLP',
'dolphin-2.8-mistral-7b': 'NLP',
'gemma-2-2b-instruct': 'NLP',
'Octopus-v2': 'NLP',
'Octopus-v4': 'NLP',
'Octo-planner': 'NLP',
'deepseek-coder-6.7b-base': 'NLP',
'Llama2-7b-chat-uncensored': 'NLP',
'Llama3-8B-Lexi-Uncensored': 'NLP',
'Llama2-7b-function-calling': 'NLP',
'lcm-dreamshaper-v7': 'Computer Vision',
'stable-diffusion-v1-5': 'Computer Vision',
'stable-diffusion-v1-4': 'Computer Vision',
'stable-diffusion-v2-1': 'Computer Vision',
'sdxl-turbo': 'Computer Vision',
'hassaku-hentai-model-v13-LCM': 'Computer Vision',
'anything-v30-LCM': 'Computer Vision',
'Phi-3-vision-128k-instruct': 'Multimodal',
'nanoLLaVA': 'Multimodal',
'llava-v1.6-mistral-7b': 'Multimodal',
'llava-v1.6-vicuna-7b': 'Multimodal',
'llava-phi-3-mini': 'Multimodal',
'llava-llama-3-8b-v1.1': 'Multimodal',
'faster-whisper-tiny.en': 'Audio',
'faster-whisper-tiny': 'Audio',
'faster-whisper-small.en': 'Audio',
'faster-whisper-small': 'Audio',
'faster-whisper-medium.en': 'Audio',
'faster-whisper-medium': 'Audio',
'faster-whisper-base.en': 'Audio',
'faster-whisper-base': 'Audio',
'faster-whisper-large-v3': 'Audio',
'whisper-tiny.en': 'Audio',
'whisper-tiny': 'Audio',
'whisper-small.en': 'Audio',
'whisper-small': 'Audio',
'whisper-base.en': 'Audio',
'whisper-base': 'Audio',
}


137 changes: 101 additions & 36 deletions nexa/general.py
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import json
import logging
from pathlib import Path

import shutil
import requests

from nexa.constants import (
Expand All @@ -13,6 +13,7 @@
NEXA_OFFICIAL_BUCKET,
NEXA_RUN_MODEL_MAP,
NEXA_TOKEN_PATH,
NEXA_OFFICIAL_MODELS_TYPE,
)


Expand Down Expand Up @@ -103,25 +104,25 @@ def pull_model(model_path):

try:
if is_model_exists(model_path):
location = get_model_location(model_path)
location, run_type = get_model_info(model_path)
logging.debug(f"Model {model_path} already exists at {location}")
return location
return location, run_type

if "/" in model_path:
success, location, model_type = pull_model_from_hub(model_path)
result = pull_model_from_hub(model_path)
else:
success, location, model_type = pull_model_from_official(model_path)
result = pull_model_from_official(model_path)

if success:
add_model_to_list(model_path, location, model_type)
logging.debug(f"Successfully pulled model {model_path} to {location}")
return location
if result["success"]:
add_model_to_list(model_path, result["local_path"], result["model_type"], result["run_type"])
logging.debug(f"Successfully pulled model {model_path} to {result['local_path']}, run_type: {result['run_type']}")
return result["local_path"], result["run_type"]
else:
logging.debug(f"Failed to pull model {model_path}")
return None
return None, "UNKNOWN"
except Exception as e:
logging.error(f"An error occurred while pulling the model: {e}")
return None
return None, "UNKNOWN"


def pull_model_from_hub(model_path):
Expand All @@ -133,13 +134,20 @@ def pull_model_from_hub(model_path):
token = file.read().strip()

try:
presigned_links = get_model_presigned_link(model_path, token)
result = get_model_presigned_link(model_path, token)
run_type = result['type']
presigned_links = result['presigned_urls']
except Exception as e:
print(f"Failed to get download models: {e}")
return False, None
return {
"success": False,
"local_path": None,
"model_type": None,
"run_type": None
}

success = True
model_location = None
local_path = None
model_type = "undefined"

# Determine model_type
Expand All @@ -159,19 +167,24 @@ def pull_model_from_hub(model_path):
download_path = NEXA_MODELS_HUB_DIR / file_path
download_file_with_progress(presigned_link, download_path)

if model_location is None:
if local_path is None:
if model_type == "onnx" or model_type == "bin":
model_location = str(download_path.parent)
local_path = str(download_path.parent)
elif model_type == "gguf":
model_location = str(download_path)
local_path = str(download_path)
else: # undefined
model_location = str(download_path.parent)
local_path = str(download_path.parent)

except Exception as e:
print(f"Failed to download {file_path}: {e}")
success = False

return success, model_location, model_type
return {
"success": success,
"local_path": local_path,
"model_type": model_type,
"run_type": run_type
}


def pull_model_from_official(model_path):
Expand All @@ -183,20 +196,33 @@ def pull_model_from_official(model_path):
model_type = "bin"
else:
model_type = "gguf"

run_type = get_run_type_from_model_path(model_path)
success, location = download_model_from_official(model_path, model_type)
return success, location, model_type

return {
"success": success,
"local_path": location,
"model_type": model_type,
"run_type": run_type
}


def get_run_type_from_model_path(model_path):
model_name, model_version = model_path.split(":")
return NEXA_OFFICIAL_MODELS_TYPE.get(model_name, "UNKNOWN")


def get_model_presigned_link(full_path, token):
"""
Get the presigned links for downloading the contents of a model folder.
Args:
full_path (str): The full path of the folder to download (e.g., "gpt2-onnx/").
token (str): The authentication token. Can be empty.
full_path (str): The full path of the folder to download (e.g., "openai/gpt2:gguf-q2_K").
token (str, optional): The authentication token. Defaults to None.
Returns:
dict: A dictionary containing the contents of the folder, where keys are file paths and values are presigned links.
dict: A dictionary containing the model type and presigned URLs.
"""

url = f"{NEXA_API_URL}/model/download-tag-folder"
Expand All @@ -205,13 +231,20 @@ def get_model_presigned_link(full_path, token):
if token:
headers["Authorization"] = f"Bearer {token}"

body = {"full_path": full_path}
body = {"full_path": full_path, "need_type": True}

try:
response = requests.post(url, headers=headers, json=body)
response.raise_for_status()
folder_contents = response.json()
return folder_contents
result = response.json()

run_type = result.get("type", [])[0] if result.get("type") else None
presigned_urls = result.get("presigned_urls", {})

return {
"run_type": run_type,
"presigned_urls": presigned_urls
}

except requests.exceptions.RequestException as e:
print(f"API request failed: {e}")
Expand Down Expand Up @@ -286,7 +319,7 @@ def is_model_exists(model_name):
return model_name in model_list


def add_model_to_list(model_name, model_location, model_type):
def add_model_to_list(model_name, model_location, model_type, run_type):
NEXA_MODEL_LIST_PATH.parent.mkdir(parents=True, exist_ok=True)

if NEXA_MODEL_LIST_PATH.exists():
Expand All @@ -295,20 +328,28 @@ def add_model_to_list(model_name, model_location, model_type):
else:
model_list = {}

model_list[model_name] = {"type": model_type, "location": model_location}
model_list[model_name] = {
"type": model_type,
"location": model_location,
"run_type": run_type
}

with open(NEXA_MODEL_LIST_PATH, "w") as f:
json.dump(model_list, f, indent=2)


def get_model_location(model_name):
def get_model_info(model_name):
if not NEXA_MODEL_LIST_PATH.exists():
return None
return None, None

with open(NEXA_MODEL_LIST_PATH, "r") as f:
model_list = json.load(f)

return model_list.get(model_name).get("location")
model_data = model_list.get(model_name, {})
location = model_data.get("location")
run_type = model_data.get("run_type")

return location, run_type


def list_models():
Expand All @@ -320,15 +361,15 @@ def list_models():
model_list = json.load(f)

table = [
(model_name, model_info["type"], model_info["location"])
(model_name, model_info["type"], model_info["run_type"], model_info["location"])
for model_name, model_info in model_list.items()
]
headers = ["Model Name", "Type", "Location"]
headers = ["Model Name", "Type", "Run Type", "Location"]
from tabulate import tabulate

print(
tabulate(
table, headers, tablefmt="pretty", colalign=("left", "left", "left")
table, headers, tablefmt="pretty", colalign=("left", "left", "left", "left")
)
)
except Exception as e:
Expand Down Expand Up @@ -359,8 +400,6 @@ def remove_model(model_path):
model_path.unlink()
print(f"Deleted model file: {model_path}")
elif model_path.is_dir():
import shutil

shutil.rmtree(model_path)
print(f"Deleted model directory: {model_path}")
else:
Expand All @@ -375,6 +414,32 @@ def remove_model(model_path):
except Exception as e:
print(f"An error occurred while removing the model: {e}")
return None

def clean():
if not NEXA_MODELS_HUB_DIR.exists():
print(f"Nothing to clean.")
return

# Ask for user confirmation
confirmation = input(f"This will remove all downloaded models and the model list. Are you sure? (y/N): ").lower().strip()

if confirmation != 'y':
print("Operation cancelled.")
return

try:
# Remove all contents of the directory
for item in NEXA_MODELS_HUB_DIR.iterdir():
if item.is_file():
item.unlink()
elif item.is_dir():
shutil.rmtree(item)

print(f"Successfully removed all contents from {NEXA_MODELS_HUB_DIR}")

except Exception as e:
print(f"An error occurred while cleaning the directory: {e}")



if __name__ == "__main__":
Expand Down
Loading

0 comments on commit c292fab

Please sign in to comment.