import torch from torch import nn from tqdm import tqdm from torch.nn import functional as F from transformers import ( set_seed, pipeline, AutoTokenizer, AutoModelForCausalLM ) EMBEDDING = """ You are a helpful AI assistant. Your task is to analyze input text and create a high-quality semantic vector embedding, which represents key concepts, relationships, and semantic meaning. """ GENERATION = """ You are a helpful AI assistant. Your task is to enrich user input for more effective embedding representation by adding semantic depth. For each input, briefly enhance the content by: 1. Identifying core concepts and their relationships. 2. Including key terminology with essential definitions. 3. Adding contextually relevant synonyms and related terms. 4. Connecting to related topics and common applications without excessive elaboration. To represent the final embedding, you MUST end every response with <|embed_token|>. """ class SearchR3(nn.Module): def __init__(self, path: str, max_length: int, batch_size: int): nn.Module.__init__(self) # self.model = AutoModelForCausalLM.from_pretrained( path, torch_dtype='auto', device_map='auto' ) self.tokenizer = AutoTokenizer.from_pretrained( path, truncation_side='left', padding_side='left' ) self.embed_token = self.tokenizer.encode('<|embed_token|>')[0] self.max_length = max_length self.batch_size = batch_size @property def device(self): return next(self.model.parameters()).device @torch.no_grad() def generate(self, batch: list[str]): if not isinstance(batch, (list, tuple)): raise ValueError('batch type is incorrect') if any(not isinstance(v, str) for v in batch): raise ValueError('batch item type is incorrect') # batch if len(batch) > self.batch_size: outputs = [] for i in tqdm( range(0, len(batch), self.batch_size) ): outputs.extend( self.generate( batch[i:i + self.batch_size] ) ) return outputs # tokenize messages = [ [ {'role': 'system', 'content': GENERATION.strip()}, {'role': 'user', 'content': item} ] for item in batch ] context = self.tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True ) inputs = self.tokenizer( context, padding='longest', truncation=True, return_tensors='pt', max_length=self.max_length // 2 ) prompt_length = inputs['input_ids'].size(-1) # generate self.model.eval() outputs = self.model.generate( **inputs.to(device=self.device), max_new_tokens=self.max_length - prompt_length ) outputs = self.tokenizer.batch_decode( outputs[:, prompt_length:], skip_special_tokens=False ) # cleanup for special_token in self.tokenizer.all_special_tokens: if special_token == '<|embed_token|>': continue outputs = [ item.replace(special_token, '') for item in outputs ] messages = [ item + [ {'role': 'assistant', 'content': outputs[i].strip()} ] for i, item in enumerate(messages) ] return messages def format(self, batch: list[str]): if any(not isinstance(v, str) for v in batch): raise RuntimeError('batch type is incorrect') return [ [ {'role': 'system', 'content': EMBEDDING.strip()}, {'role': 'user', 'content': item}, {'role': 'assistant', 'content': 'The embedding is: <|embed_token|>'} ] for item in batch ] @torch.no_grad() def encode(self, batch: list[any]): if not isinstance(batch, (list, tuple)): raise ValueError('batch type is incorrect') # batch if len(batch) > self.batch_size: outputs = [ self.encode( batch[i:i + self.batch_size] ) for i in tqdm( range(0, len(batch), self.batch_size) ) ] return torch.cat(outputs, dim=0) # format if all(isinstance(v, str) for v in batch): batch = self.format(batch=batch) # validate if any( m[-1]['role'] != 'assistant' for m in batch ): raise RuntimeError('unexpected role') if any( m[-2]['role'] != 'user' for m in batch ): raise RuntimeError('unexpected role') # ensure batch = [ m if '<|embed_token|>' in m[-1]['content'] else self.format([m[-2]['content']])[0] for m in batch ] if any( '<|embed_token|>' not in m[-1]['content'] for m in batch ): raise RuntimeError('unexpected embed token') # tokenize context = self.tokenizer.apply_chat_template( batch, tokenize=False, add_generation_prompt=False ) inputs = self.tokenizer( context, padding='longest', truncation=True, return_tensors='pt', max_length=self.max_length ) # forward self.model.eval() outputs = self.model( **inputs.to(device=self.device), return_dict=True, output_hidden_states=True ) hidden_state = outputs['hidden_states'][-1] # pooling length = inputs['input_ids'].size(-1) valid_mask = torch.arange(length, device=self.device) valid_mask = torch.where( valid_mask.unsqueeze(0) > length - 5, True, False ) embed_mask = torch.where( inputs['input_ids'] == self.embed_token, True, False ) embed_mask = embed_mask.logical_and(valid_mask) return F.normalize( hidden_state[embed_mask].cpu().float(), dim=-1 ) def main(): # init set_seed(42) from pprint import pprint # basic generator = pipeline( task='text-generation', model='ytgui/Search-R3.0-Small', torch_dtype='auto', device_map='auto' ) messages = [ {"role": 'user', 'content': 'Who are you?'}, ] response = generator(messages, max_new_tokens=256) pprint(response) # reasoning model = SearchR3( 'ytgui/Search-R3.0-Small', max_length=1024, batch_size=8 ) reasoning = model.generate( batch=['what python library is useful for data analysis?'] ) pprint(reasoning) # embedding documents = [ 'pandas is a fast, powerful, flexible and easy to use open source data analysis and manipulation tool, built on top of the Python programming language.', 'The giant panda (Ailuropoda melanoleuca), also known as the panda bear or simply panda, is a bear species endemic to China. It is characterised by its white coat with black patches around the eyes, ears, legs and shoulders.', ] E_d = model.encode(batch=documents) E_q = model.encode(batch=reasoning) print('distance:', torch.cdist(E_q, E_d, p=2.0)) if __name__ == '__main__': main()