nodes/Big-Idea-Technology-ComfyUI_LLM_Node
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

add llama and mistral model support

Browse Source
This commit is contained in:
nazgut
2024-03-29 08:06:45 +01:00
committed by GitHub
parent 3ccea4b90a
commit 2e60cf3778
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
LLM_Node.py
+2 -2
View File
@@ -111,7 +111,7 @@ class LLM_Node:
# Dynamically loading the model based on its type
if config.model_type == "t5":
model = AutoModelForSeq2SeqLM.from_pretrained(model_path, **model_kwargs)
elif config.model_type in ["gpt2", "gpt_refact", "gemma"]:
elif config.model_type in ["gpt2", "gpt_refact", "gemma", "llama", "mistral"]:
model = AutoModelForCausalLM.from_pretrained(model_path, **model_kwargs)
elif config.model_type == "bert":
model = AutoModelForSequenceClassification.from_pretrained(model_path, **model_kwargs)
@@ -127,7 +127,7 @@ class LLM_Node:
if option in AdvOptionsConfig:
generate_kwargs[option] = AdvOptionsConfig[option]
if config.model_type in ["t5", "gpt2", "gpt_refact", "gemma"]:
if config.model_type in ["t5", "gpt2", "gpt_refact", "gemma", "llama", "mistral"]:
input_ids = tokenizer(text, return_tensors="pt").input_ids.to(self.device)
outputs = model.generate(input_ids, **generate_kwargs)
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)