feat: create deployment scripts

This commit is contained in:
nobody 2025-11-02 13:09:23 -08:00
commit 8d5bce4bfb
Signed by: GrocerPublishAgent
GPG key ID: D460CD54A9E3AB86
22 changed files with 2697 additions and 74 deletions

View file

@ -8,6 +8,7 @@ import os
# Set NLTK data path to project directory
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)
@ -28,12 +29,14 @@ AVAILABLE_MODELS = {
}
# On clustering
# all-mpnet-base-v2: 40.03
# mixedbread-ai/mxbai-embed-large-v1: 46.71
# gte-large-en-v1.5: 47.95
# Qwen/Qwen3-Embedding-0.6B: 52.33
# Qwen/Qwen3-Embedding-4B: 57.15
# On STS
# all-mpnet-base-v2: 58.17
# gte-large-en-v1.5: 81.43
# Qwen/Qwen3-Embedding-0.6B: 76.17
# Qwen/Qwen3-Embedding-4B: 80.86
@ -43,23 +46,22 @@ AVAILABLE_MODELS = {
print("Loading sentence transformer models...")
models = {}
models['all-mpnet-base-v2'] = SentenceTransformer('all-mpnet-base-v2')
models['all-mpnet-base-v2'] = SentenceTransformer('all-mpnet-base-v2', cache_folder=TRANSFORMERS_CACHE_DIR)
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)
models['gte-large-en-v1.5'] = SentenceTransformer('Alibaba-NLP/gte-large-en-v1.5', trust_remote_code=True, cache_folder=TRANSFORMERS_CACHE_DIR)
#print("Loading Qwen/Qwen3-Embedding-4B")
#models['qwen3-embedding-4b'] = SentenceTransformer('Qwen/Qwen3-Embedding-4B', trust_remote_code=True)
#models['qwen3-embedding-4b'] = SentenceTransformer('Qwen/Qwen3-Embedding-4B', trust_remote_code=True, cache_folder=TRANSFORMERS_CACHE_DIR)
print("Loading mixedbread-ai/mxbai-embed-large-v1")
models["mxbai-embed-large-v1"] = SentenceTransformer('mixedbread-ai/mxbai-embed-large-v1')
models["mxbai-embed-large-v1"] = SentenceTransformer('mixedbread-ai/mxbai-embed-large-v1', cache_folder=TRANSFORMERS_CACHE_DIR)
print("All models loaded!")
sent_detector = nltk.data.load('tokenizers/punkt/english.pickle')
def cos_sim(a, b):
sims = a @ b.T
a_norm = np.linalg.norm(a, axis=-1)
b_norm = np.linalg.norm(b, axis=-1)
a_normalized = (sims.T / a_norm.T).T
sims = a_normalized / b_norm
def cos_sim(a):
sims = a @ a.T
a_norm = np.linalg.norm(a, axis=-1, keepdims=True)
sims /= a_norm
sims /= a_norm.T
return sims
def degree_power(A, k):
@ -79,7 +81,7 @@ def get_sentences(source_text):
def text_rank(sentences, model_name='all-mpnet-base-v2'):
model = models[model_name]
vectors = model.encode(sentences)
adjacency = torch.tensor(cos_sim(vectors, vectors)).fill_diagonal_(0.)
adjacency = torch.tensor(cos_sim(vectors)).fill_diagonal_(0.)
adjacency[adjacency < 0] = 0
return normalized_adjacency(adjacency)