From 6ddc63981c7ff4ef09639d158b4632ce679cfbe1 Mon Sep 17 00:00:00 2001 From: Partho Das Date: Thu, 14 Dec 2023 20:26:27 +0530 Subject: [PATCH] mistral loading fix --- src/llm_vm/onsite_llm.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/llm_vm/onsite_llm.py b/src/llm_vm/onsite_llm.py index 05b8c05a..3fee8763 100644 --- a/src/llm_vm/onsite_llm.py +++ b/src/llm_vm/onsite_llm.py @@ -12,6 +12,7 @@ GPTNeoForCausalLM, GPTNeoXForCausalLM, LlamaForCausalLM, + MistralForCausalLM, LlamaTokenizer, CodeLlamaTokenizer, DataCollatorForLanguageModeling, @@ -533,7 +534,7 @@ class SmallLocalOpenMistral(BaseOnsiteLLM): model_uri="Open-Orca/Mistral-7B-OpenOrca" def model_loader(self): - return LlamaForCausalLM.from_pretrained(self.model_uri) + return MistralForCausalLM.from_pretrained(self.model_uri) def tokenizer_loader(self): return LlamaTokenizer.from_pretrained(self.model_uri)