oss/salience-editor
oss/salience-editor
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: deploy model api server to chicago-web01

Browse source
This commit is contained in:
nobody 2025-11-27 11:01:54 -08:00
commit 0cb89ddc80
Signed by: GrocerPublishAgent
GPG key ID: D460CD54A9E3AB86
6 changed files with 394 additions and 18 deletions
Download patch file Download diff file Expand all files Collapse all files

42
api/salience/salience.py
View file

@ -1,24 +1,34 @@
import numpy as np
import torch
from sentence_transformers import SentenceTransformer
import nltk.data
import nltk
import os
# Set NLTK data path to project directory
# Set default cache locations BEFORE importing libraries that use them
PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
NLTK_DATA_DIR = os.path.join(PROJECT_DIR, 'nltk_data')
TRANSFORMERS_CACHE_DIR = os.path.join(PROJECT_DIR, 'models_cache')
# Add to NLTK's search path
nltk.data.path.insert(0, NLTK_DATA_DIR)
if 'NLTK_DATA' not in os.environ:
nltk_data_path = os.path.join(PROJECT_DIR, 'cache-nltk')
os.makedirs(nltk_data_path, exist_ok=True)
os.environ['NLTK_DATA'] = nltk_data_path
# Download to the custom location
if 'HF_HOME' not in os.environ:
os.environ['HF_HOME'] = os.path.join(PROJECT_DIR, 'cache-huggingface')
from salience.timed_import import timed_import
with timed_import("import numpy as np"):
import numpy as np
with timed_import("import torch"):
import torch
with timed_import("from sentence_transformers import SentenceTransformer"):
from sentence_transformers import SentenceTransformer
with timed_import("import nltk"):
import nltk.data
import nltk
# Download punkt_tab to the configured location
# Using punkt_tab (the modern tab-separated format introduced in NLTK 3.8+)
# instead of the older punkt pickle format
# The punkt_tab model version depends on the NLTK Python package version
# Check your NLTK version with: uv pip show nltk
nltk.download('punkt_tab', download_dir=NLTK_DATA_DIR)
nltk.download('punkt_tab')
# Available models for the demo
AVAILABLE_MODELS = {
@ -46,13 +56,13 @@ AVAILABLE_MODELS = {
print("Loading sentence transformer models...")
models = {}
models['all-mpnet-base-v2'] = SentenceTransformer('all-mpnet-base-v2', cache_folder=TRANSFORMERS_CACHE_DIR)
models['all-mpnet-base-v2'] = SentenceTransformer('all-mpnet-base-v2')
print("Loading Alibaba-NLP/gte-large-en-v1.5")
models['gte-large-en-v1.5'] = SentenceTransformer('Alibaba-NLP/gte-large-en-v1.5', trust_remote_code=True, cache_folder=TRANSFORMERS_CACHE_DIR)
models['gte-large-en-v1.5'] = SentenceTransformer('Alibaba-NLP/gte-large-en-v1.5', trust_remote_code=True)
#print("Loading Qwen/Qwen3-Embedding-4B")
#models['qwen3-embedding-4b'] = SentenceTransformer('Qwen/Qwen3-Embedding-4B', trust_remote_code=True, cache_folder=TRANSFORMERS_CACHE_DIR)
#models['qwen3-embedding-4b'] = SentenceTransformer('Qwen/Qwen3-Embedding-4B', trust_remote_code=True)
print("Loading mixedbread-ai/mxbai-embed-large-v1")
models["mxbai-embed-large-v1"] = SentenceTransformer('mixedbread-ai/mxbai-embed-large-v1', cache_folder=TRANSFORMERS_CACHE_DIR)
models["mxbai-embed-large-v1"] = SentenceTransformer('mixedbread-ai/mxbai-embed-large-v1')
print("All models loaded!")
sent_detector = nltk.data.load('tokenizers/punkt/english.pickle')