@@ -543,18 +543,24 @@ def __init__(self, tokenizer_path=None, max_length=77, pad_with_end=True, embedd
543543 def _try_get_embedding (self , embedding_name :str ):
544544 '''
545545 Takes a potential embedding name and tries to retrieve it.
546- Returns a Tuple consisting of the embedding and any leftover string, embedding can be None.
546+ Returns a Tuple consisting of the embedding, the cleaned embedding name, and any leftover string, embedding can be None.
547547 '''
548548 split_embed = embedding_name .split ()
549549 embedding_name = split_embed [0 ]
550550 leftover = ' ' .join (split_embed [1 :])
551+
552+ match = re .search (r'[<\[]' , embedding_name )
553+ if match is not None :
554+ leftover = embedding_name [match .start ():] + (" " + leftover if leftover else "" )
555+ embedding_name = embedding_name [:match .start ()]
556+
551557 embed = load_embed (embedding_name , self .embedding_directory , self .embedding_size , self .embedding_key )
552558 if embed is None :
553559 stripped = embedding_name .strip (',' )
554560 if len (stripped ) < len (embedding_name ):
555561 embed = load_embed (stripped , self .embedding_directory , self .embedding_size , self .embedding_key )
556- return (embed , "{} {}" .format (embedding_name [len (stripped ):], leftover ))
557- return (embed , leftover )
562+ return (embed , embedding_name , "{} {}" .format (embedding_name [len (stripped ):], leftover ))
563+ return (embed , embedding_name , leftover )
558564
559565 def pad_tokens (self , tokens , amount ):
560566 if self .pad_left :
@@ -585,7 +591,7 @@ def tokenize_with_weights(self, text:str, return_word_ids=False, tokenizer_optio
585591 tokens = []
586592 for weighted_segment , weight in parsed_weights :
587593 to_tokenize = unescape_important (weighted_segment )
588- split = re .split (' {0}| \n {0 }' .format (self .embedding_identifier ), to_tokenize )
594+ split = re .split (r'(?<=\s){ }' .format (re . escape ( self .embedding_identifier ) ), to_tokenize )
589595 to_tokenize = [split [0 ]]
590596 for i in range (1 , len (split )):
591597 to_tokenize .append ("{}{}" .format (self .embedding_identifier , split [i ]))
@@ -595,7 +601,7 @@ def tokenize_with_weights(self, text:str, return_word_ids=False, tokenizer_optio
595601 # if we find an embedding, deal with the embedding
596602 if word .startswith (self .embedding_identifier ) and self .embedding_directory is not None :
597603 embedding_name = word [len (self .embedding_identifier ):].strip ('\n ' )
598- embed , leftover = self ._try_get_embedding (embedding_name )
604+ embed , embedding_name , leftover = self ._try_get_embedding (embedding_name )
599605 if embed is None :
600606 logging .warning (f"warning, embedding:{ embedding_name } does not exist, ignoring" )
601607 else :
0 commit comments