Update app.py
Browse files
app.py
CHANGED
@@ -8,8 +8,9 @@ import queue
|
|
8 |
import torch
|
9 |
|
10 |
# Загружаем модель
|
11 |
-
model_name = "HIT-TMG/KaLM-embedding-multilingual-mini-instruct-v1"
|
12 |
model = SentenceTransformer(model_name)
|
|
|
13 |
|
14 |
# Имя файла для сохранения эмбеддингов
|
15 |
embeddings_file = f"movie_embeddings_{model_name.replace('/', '_')}.json"
|
@@ -55,11 +56,17 @@ search_in_progress = False
|
|
55 |
movie_embeddings_lock = threading.Lock()
|
56 |
|
57 |
# Размер пакета для обработки эмбеддингов
|
58 |
-
batch_size =
|
59 |
|
60 |
-
|
61 |
-
|
62 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
63 |
|
64 |
def process_movies():
|
65 |
"""
|
@@ -91,7 +98,7 @@ def process_movies():
|
|
91 |
]
|
92 |
|
93 |
print(f"Создаются эмбеддинги для фильмов: {', '.join(titles)}...")
|
94 |
-
embeddings = model.encode(embedding_strings, convert_to_tensor=True, batch_size=batch_size).tolist()
|
95 |
|
96 |
with movie_embeddings_lock:
|
97 |
for title, embedding in zip(titles, embeddings):
|
@@ -105,7 +112,8 @@ def process_movies():
|
|
105 |
|
106 |
def get_query_embedding(query):
|
107 |
"""
|
108 |
-
Возвращает эмбеддинг для
|
|
|
109 |
Иначе создает эмбеддинг, сохраняет его и возвращает.
|
110 |
"""
|
111 |
if query in query_embeddings:
|
@@ -113,7 +121,7 @@ def get_query_embedding(query):
|
|
113 |
return query_embeddings[query]
|
114 |
else:
|
115 |
print(f"Создается эмбеддинг для запроса '{query}'...")
|
116 |
-
embedding = encode_string(query).tolist()
|
117 |
query_embeddings[query] = embedding
|
118 |
# Сохраняем эмбеддинги запросов в файл
|
119 |
with open(query_embeddings_file, "w", encoding="utf-8") as f:
|
@@ -123,7 +131,7 @@ def get_query_embedding(query):
|
|
123 |
|
124 |
def search_movies(query, top_k=10):
|
125 |
"""
|
126 |
-
Ищет наиболее похожие фильмы по
|
127 |
|
128 |
Args:
|
129 |
query: Текстовый запрос.
|
@@ -138,7 +146,7 @@ def search_movies(query, top_k=10):
|
|
138 |
print(f"\n\033[1mПоиск по запросу: '{query}'\033[0m")
|
139 |
|
140 |
print(f"Начало создания эмбеддинга для запроса: {time.strftime('%Y-%m-%d %H:%M:%S')}")
|
141 |
-
query_embedding_tensor = encode_string(query)
|
142 |
print(f"Окончание создания эмбеддинга для запроса: {time.strftime('%Y-%m-%d %H:%M:%S')}")
|
143 |
|
144 |
with movie_embeddings_lock:
|
|
|
8 |
import torch
|
9 |
|
10 |
# Загружаем модель
|
11 |
+
model_name = "HIT-TMG/KaLM-embedding-multilingual-mini-instruct-v1.5"
|
12 |
model = SentenceTransformer(model_name)
|
13 |
+
# model.max_seq_length = 8192 # Убираем явное ограничение длины последовательности
|
14 |
|
15 |
# Имя файла для сохранения эмбеддингов
|
16 |
embeddings_file = f"movie_embeddings_{model_name.replace('/', '_')}.json"
|
|
|
56 |
movie_embeddings_lock = threading.Lock()
|
57 |
|
58 |
# Размер пакета для обработки эмбеддингов
|
59 |
+
batch_size = 32 # Увеличиваем размер пакета в 2 раза
|
60 |
|
61 |
+
# Инструкция для запроса
|
62 |
+
query_prompt = "Инструкция: Найди релевантные фильмы по запросу. \n Запрос: "
|
63 |
+
|
64 |
+
def encode_string(text, prompt=None):
|
65 |
+
"""Кодирует строку в эмбеддинг с использованием инструкции, если она задана."""
|
66 |
+
if prompt:
|
67 |
+
return model.encode(text, prompt=prompt, convert_to_tensor=True, normalize_embeddings=True)
|
68 |
+
else:
|
69 |
+
return model.encode(text, convert_to_tensor=True, normalize_embeddings=True)
|
70 |
|
71 |
def process_movies():
|
72 |
"""
|
|
|
98 |
]
|
99 |
|
100 |
print(f"Создаются эмбеддинги для фильмов: {', '.join(titles)}...")
|
101 |
+
embeddings = model.encode(embedding_strings, convert_to_tensor=True, batch_size=batch_size, normalize_embeddings=True).tolist()
|
102 |
|
103 |
with movie_embeddings_lock:
|
104 |
for title, embedding in zip(titles, embeddings):
|
|
|
112 |
|
113 |
def get_query_embedding(query):
|
114 |
"""
|
115 |
+
Возвращает эмбеддинг для запроса с инструкцией.
|
116 |
+
Если эмбеддинг уже создан, возвращает его из словаря.
|
117 |
Иначе создает эмбеддинг, сохраняет его и возвращает.
|
118 |
"""
|
119 |
if query in query_embeddings:
|
|
|
121 |
return query_embeddings[query]
|
122 |
else:
|
123 |
print(f"Создается эмбеддинг для запроса '{query}'...")
|
124 |
+
embedding = encode_string(query, prompt=query_prompt).tolist()
|
125 |
query_embeddings[query] = embedding
|
126 |
# Сохраняем эмбеддинги запросов в файл
|
127 |
with open(query_embeddings_file, "w", encoding="utf-8") as f:
|
|
|
131 |
|
132 |
def search_movies(query, top_k=10):
|
133 |
"""
|
134 |
+
Ищет наиболее похожие фильмы по запросу с использованием инструкции.
|
135 |
|
136 |
Args:
|
137 |
query: Текстовый запрос.
|
|
|
146 |
print(f"\n\033[1mПоиск по запросу: '{query}'\033[0m")
|
147 |
|
148 |
print(f"Начало создания эмбеддинга для запроса: {time.strftime('%Y-%m-%d %H:%M:%S')}")
|
149 |
+
query_embedding_tensor = encode_string(query, prompt=query_prompt)
|
150 |
print(f"Окончание создания эмбеддинга для запроса: {time.strftime('%Y-%m-%d %H:%M:%S')}")
|
151 |
|
152 |
with movie_embeddings_lock:
|