From 2e60cf377801ded14cdbd7ae5388d2c7e404b556 Mon Sep 17 00:00:00 2001 From: nazgut Date: Fri, 29 Mar 2024 08:06:45 +0100 Subject: [PATCH] add llama and mistral model support --- LLM_Node.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/LLM_Node.py b/LLM_Node.py index 440a0a9..f1f2683 100644 --- a/LLM_Node.py +++ b/LLM_Node.py @@ -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)