opex792 commited on
Commit
429c3b2
·
verified ·
1 Parent(s): 81516d6

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +18 -10
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 = 16 # Подберите оптимальный размер
59
 
60
- def encode_string(text):
61
- """Кодирует строку в эмбеддинг."""
62
- return model.encode(text, convert_to_tensor=True)
 
 
 
 
 
 
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: