ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-03-11 21:49:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

Bug fixes and improvements to token injection

Browse Source
This commit is contained in:
Jaret Burkett
2023-09-08 06:10:59 -06:00
parent 92a086d5a5
commit ce4f9fe02a
5 changed files with 74 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
toolkit/embedding.py
View File

@@ -141,8 +141,8 @@ class Embedding:
# however, on training we don't use that pipeline, so we have to do it ourselves
def inject_embedding_to_prompt(self, prompt, expand_token=False, to_replace_list=None, add_if_not_present=True):
output_prompt = prompt
embedding_tokens = self.embedding_tokens[0] # shoudl be the same
default_replacements = [self.name, self.trigger, "[name]", "[trigger]", embedding_tokens]
embedding_tokens = self.embedding_tokens[0] # shoudl be the same
default_replacements = ["[name]", "[trigger]"]
replace_with = embedding_tokens if expand_token else self.trigger
if to_replace_list is None:
@@ -167,7 +167,7 @@ class Embedding:
if num_instances > 1:
print(
f"Warning: {self.name} token appears {num_instances} times in prompt {output_prompt}. This may cause issues.")
f"Warning: {replace_with} token appears {num_instances} times in prompt {output_prompt}. This may cause issues.")
return output_prompt