Spaces:
Running
Running
| # routers/embedding/__init__.py | |
| import os | |
| import re | |
| import sys | |
| import threading | |
| import torch | |
| from sentence_transformers import SentenceTransformer, util | |
| from typing import Dict, List, Tuple, Set, LiteralString | |
| class EmbeddingContext: | |
| # These don't change | |
| TOKEN_LEN_MAX_FOR_EMBEDDING = 512 | |
| # Set when creating the object | |
| lock = None | |
| model = None | |
| openai_client = None | |
| model_name = '' | |
| config_type = '' | |
| embedding_shape = None | |
| embedding_dtype = None | |
| embedding_device = None | |
| # Updates constantly | |
| data = {} | |
| def __init__(self): | |
| try: | |
| from config import settings | |
| except: | |
| sys.path.append(os.path.abspath( | |
| os.path.join(os.path.dirname(__file__), '../..'))) | |
| from config import settings | |
| self.lock = threading.Lock() | |
| config_type = settings.embedding_api | |
| model_name = settings.embedding_model | |
| if config_type == 'sbert': | |
| self.model = SentenceTransformer(model_name, use_auth_token=False) | |
| self.model.max_seq_length = self.TOKEN_LEN_MAX_FOR_EMBEDDING | |
| print("Max Sequence Length:", self.model.max_seq_length) | |
| self.encode = self.encode_sbert | |
| if torch.cuda.is_available(): | |
| self.model = self.model.to('cuda') | |
| elif config_type == 'openai': | |
| from openai import OpenAI | |
| self.openai_client = OpenAI( | |
| # base_url = settings.openai_api_base | |
| api_key=settings.OPENAI_API_KEY, | |
| ) | |
| self.encode = self.encode_openai | |
| self.model_name = model_name | |
| self.config_type = config_type | |
| tmp = self.encode(['tmp']) | |
| self.embedding_shape = tmp.shape[1:] | |
| self.embedding_dtype = tmp.dtype | |
| self.embedding_device = tmp.device | |
| def encode(self, texts_to_embed): | |
| pass | |
| def encode_sbert(self, texts_to_embed): | |
| return self.model.encode(texts_to_embed, show_progress_bar=True, convert_to_tensor=True, normalize_embeddings=True) | |
| def encode_openai(self, texts_to_embed): | |
| import math | |
| import time | |
| tokens_count = 0 | |
| for text in texts_to_embed: | |
| tokens_count += len(self.get_tokens(text)) | |
| chunks_num = math.ceil(tokens_count / 500000) | |
| chunk_size = math.ceil(len(texts_to_embed) / chunks_num) | |
| embeddings = [] | |
| for i in range(chunks_num): | |
| start = i * chunk_size | |
| end = start + chunk_size | |
| chunk = texts_to_embed[start:end] | |
| embeddings_tmp = self.openai_client.embeddings.create( | |
| model=self.model_name, | |
| input=chunk, | |
| ).data | |
| if embeddings_tmp is None: | |
| break | |
| embeddings.extend(embeddings_tmp) | |
| if i < chunks_num - 1: | |
| time.sleep(60) # Wait 1 minute before the next call | |
| return torch.stack([torch.tensor(embedding.embedding, dtype=torch.float32) for embedding in embeddings]) | |
| def get_tokens(self, text): | |
| if self.model: | |
| return self.model.tokenizer.tokenize(text) | |
| tokens = [] | |
| for token in re.split(r'(\W|\b)', text): | |
| if token.strip(): | |
| tokens.append(token) | |
| return tokens | |
| class SplitDocs: | |
| def split_in_topics(self, | |
| filedir: LiteralString = None, | |
| *, | |
| pattern_filename=r'(?<!navigation)\.(md|rst)', | |
| pattern_content_sub=r'---\nhide:[\s\S]+?---\s*', | |
| patterns_titles=( | |
| r'^# (.+)', r'^## (.+)', r'^### (.+)'), | |
| ) -> List[Tuple[str, str]]: | |
| def matches_pattern(filename): | |
| return re.search(pattern_filename, filename) is not None | |
| def split_patterns_recursive(patterns, text, index=-1): | |
| sections = re.split(patterns[0], text, flags=re.MULTILINE) | |
| for i, section in enumerate(sections): | |
| if not section.strip(): | |
| continue | |
| is_match = bool(i & 1) | |
| if is_match: | |
| yield (index, section) | |
| elif len(patterns) > 1: | |
| for j, section_j in split_patterns_recursive(patterns[1:], section, index + 1): | |
| yield (j, section_j) | |
| else: | |
| yield (-1, section) | |
| for root, _, files in os.walk(filedir): | |
| for name in files: | |
| if not matches_pattern(name): | |
| continue | |
| full_path = os.path.join(root, name) | |
| with open(full_path, 'r', encoding='utf-8') as file: | |
| content = file.read() | |
| if pattern_content_sub: | |
| content = re.sub(pattern_content_sub, '', content) | |
| rel_path = full_path.replace(filedir, '').replace('\\', '/') | |
| # Protect code parts | |
| patterns = (r'(```[\s\S]+?```)', *patterns_titles) | |
| last_titles = [] | |
| last_titles_index = [] | |
| content_accum = '' | |
| for i, section in split_patterns_recursive(patterns, content): | |
| if i < 0: | |
| content_accum += section | |
| continue | |
| if content_accum: | |
| yield rel_path, last_titles, content_accum | |
| content_accum = '' | |
| if not last_titles_index or i > last_titles_index[-1]: | |
| last_titles_index.append(i) | |
| last_titles.append(section) | |
| continue | |
| while len(last_titles_index) > 1 and i < last_titles_index[-1]: | |
| last_titles_index.pop() | |
| last_titles.pop() | |
| # Replace | |
| last_titles_index[-1] = i | |
| last_titles[-1] = section | |
| if content_accum or i != -1: | |
| yield rel_path, last_titles, content_accum | |
| def reduce_text(_self, text): | |
| text = re.sub(r'^\n+', '', text) # Strip | |
| text = re.sub(r'<.*?>', '', text) # Remove HTML tags | |
| text = re.sub(r':\S*: ', '', text) # Remove [:...:] patterns | |
| text = re.sub(r'\s*\n+', '\n', text) | |
| return text | |
| def embedding_header(_self, rel_path, titles): | |
| return f"{rel_path}\n# {' | '.join(titles)}\n\n" | |
| def split_for_embedding(self, | |
| filedir: LiteralString = None, | |
| *, | |
| pattern_filename=r'(?<!navigation)\.(md|rst)', | |
| pattern_content_sub=r'---\nhide:[\s\S]+?---\s*', | |
| patterns_titles=( | |
| r'^# (.+)', r'^## (.+)', r'^### (.+)'), | |
| ): | |
| tokenizer = EMBEDDING_CTX.model.tokenizer | |
| max_tokens = EMBEDDING_CTX.model.max_seq_length | |
| texts = [] | |
| for rel_path, titles, content in self.split_in_topics( | |
| filedir, pattern_filename=pattern_filename, pattern_content_sub=pattern_content_sub, patterns_titles=patterns_titles): | |
| header = self.embedding_header(rel_path, titles) | |
| tokens_pre_len = len(tokenizer.tokenize(header)) | |
| tokens_so_far = tokens_pre_len | |
| text_so_far = header | |
| for part in self.reduce_text(content).splitlines(): | |
| part += '\n' | |
| part_tokens_len = len(tokenizer.tokenize(part)) | |
| if tokens_so_far + part_tokens_len > max_tokens: | |
| texts.append(text_so_far) | |
| text_so_far = header | |
| tokens_so_far = tokens_pre_len | |
| text_so_far += part | |
| tokens_so_far += part_tokens_len | |
| if tokens_so_far != tokens_pre_len: | |
| texts.append(text_so_far) | |
| return texts | |
| EMBEDDING_CTX = EmbeddingContext() | |