Update pipeline.py

pipeline.py

@@ -2,59 +2,32 @@ import torch
 import transformers
 from accelerate import dispatch_model, infer_auto_device_map
 from accelerate.utils import get_balanced_memory
+from transformers import BitsAndBytesConfig, StoppingCriteria, StoppingCriteriaList
 from typing import Dict, List, Any
 
 class PreTrainedPipeline():
     def __init__(self, path=""):
         path = "oleksandrfluxon/mpt-7b-instruct-evaluate"
         print("===> path", path)
-        
-        with torch.autocast('cuda'):
-            config = transformers.AutoConfig.from_pretrained(
-                path, 
-                trust_remote_code=True
-            )
-            # config.attn_config['attn_impl'] = 'triton'
-            config.init_device = 'cuda:0' # For fast initialization directly on GPU!
-            config.max_seq_len = 4096 # (input + output) tokens can now be up to 4096
-    
-            print("===> loading model")
-            model = transformers.AutoModelForCausalLM.from_pretrained(
-              path,
-              config=config,
-              # torch_dtype=torch.bfloat16, # Load model weights in bfloat16
-              torch_dtype=torch.float16,
-              trust_remote_code=True,
-              device_map={"": 0},
-              load_in_8bit=True # Load model in the lowest 4-bit precision quantization
-            )
-            # model.to('cuda')
-            print("===> model loaded")
-    
-            # removed device_map="auto"
-            tokenizer = transformers.AutoTokenizer.from_pretrained('EleutherAI/gpt-neox-20b', padding_side="left")
 
+        device = f'cuda:{cuda.current_device()}' if cuda.is_available() else 'cpu'
+        print("===> device", device)
 
-            max_memory = get_balanced_memory(
-                model,
-                max_memory=None,
-                no_split_module_classes=["MPTBlock"],
-                dtype='float16',
-                low_zero=False
-            )
-            
-            device_map = infer_auto_device_map(
-                model,
-                max_memory=max_memory,
-                no_split_module_classes=["MPTBlock"],
-                dtype='float16'
-            )
-            model = dispatch_model(model, device_map=device_map)
-                
+        model = transformers.AutoModelForCausalLM.from_pretrained(
+            'oleksandrfluxon/mpt-7b-instruct-evaluate',
+            trust_remote_code=True,
+            load_in_8bit=True,  # this requires the `bitsandbytes` library
+            max_seq_len=8192,
+            init_device=device
+        )
+        model.eval()
+        #model.to(device)
+        print(f"===> Model loaded on {device}")
 
-            # device='cuda:0'
-            self.pipeline = transformers.pipeline('text-generation', model=model, tokenizer=tokenizer)
-            print("===> init finished")
+        tokenizer = transformers.AutoTokenizer.from_pretrained("mosaicml/mpt-7b")
+
+        self.pipeline = transformers.pipeline('text-generation', model=model, tokenizer=tokenizer)
+        print("===> init finished")
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
@@ -71,8 +44,7 @@ class PreTrainedPipeline():
         print("===> inputs", inputs)
         print("===> parameters", parameters)
 
-        with torch.autocast('cuda'):
-            result = self.pipeline(inputs, **parameters)
-            print("===> result", result)
+        result = self.pipeline(inputs, **parameters)
+        print("===> result", result)
 
         return result