update

requirements.txt

@@ -13,3 +13,4 @@ torch==2.0.0
 transformers @ git+https://github.com/huggingface/transformers@ef42c2c487260c2a0111fa9d17f2507d84ddedea
 unstructured==0.6.2
 xformers==0.0.19
+ctransformers

training/generate.py

@@ -1,10 +1,11 @@
 import logging
 import re
 from typing import List, Tuple
+from ctransformers import AutoModelForCausalLM
 
 import numpy as np
 from transformers import (
-    AutoModelForCausalLM,
+    # AutoModelForCausalLM,
     AutoTokenizer,
     Pipeline,
     PreTrainedModel,
@@ -32,10 +33,14 @@ def load_model_tokenizer_for_generate(
     Returns:
         Tuple[PreTrainedModel, PreTrainedTokenizer]: model and tokenizer
     """
-    tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path, padding_side="left", cache_dir="/media/siiva/DataStore/LLMs/cache/dollyV2")
-    model = AutoModelForCausalLM.from_pretrained(
-        pretrained_model_name_or_path, device_map="auto", trust_remote_code=True, cache_dir="/media/siiva/DataStore/LLMs/cache/dollyV2"
-    )
+    # tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path, padding_side="left")#, cache_dir="/media/siiva/DataStore/LLMs/cache/dollyV2")
+    # model = AutoModelForCausalLM.from_pretrained(
+    #     pretrained_model_name_or_path, device_map="auto", trust_remote_code=True)#, cache_dir="/media/siiva/DataStore/LLMs/cache/dollyV2"
+    #)
+    tokenizer = AutoTokenizer.from_pretrained('dolly-v2')
+
+    model = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path, model_type='dolly-v2')
+
     return model, tokenizer