Create pipeline.py

pipeline.py

@@ -0,0 +1,78 @@
+import torch
+import transformers
+from accelerate import dispatch_model, infer_auto_device_map
+from accelerate.utils import get_balanced_memory
+from typing import Dict, List, Any
+
+class PreTrainedPipeline():
+    def __init__(self, path=""):
+        path = "oleksandrfluxon/mpt-7b-chat-4bit"
+        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="auto", 
+              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")
+
+
+            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)
+                
+
+            # device='cuda:0'
+            self.pipeline = transformers.pipeline('text-generation', model=model, tokenizer=tokenizer)
+            print("===> init finished")
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+       data args:
+            inputs (:obj: `str`)
+            parameters (:obj: `str`)
+      Return:
+            A :obj:`str`: todo
+        """
+        # get inputs
+        inputs = data.pop("inputs",data)
+        parameters = data.pop("parameters", {})
+        date = data.pop("date", None)
+        print("===> inputs", inputs)
+        print("===> parameters", parameters)
+
+        with torch.autocast('cuda'):
+            result = self.pipeline(inputs, **parameters)
+            print("===> result", result)
+
+        return result