Create pipeline.py

pipeline.py

@@ -0,0 +1,46 @@
+import torch
+import transformers
+from typing import Dict, List, Any
+
+class PreTrainedPipeline():
+    def __init__(self, path=""):
+        path = "oleksandrfluxon/mpt-7b-instruct-2"
+        print("===> path", path)
+        
+        config = transformers.AutoConfig.from_pretrained(name, trust_remote_code=True)
+        config.max_seq_len = 4096 # (input + output) tokens can now be up to 4096
+
+        print("===> loading model")
+        model = transformers.AutoModelForCausalLM.from_pretrained(
+          name,
+          config=config,
+          torch_dtype=torch.bfloat16, # Load model weights in bfloat16
+          trust_remote_code=True,
+          load_in_4bit=True, # Load model in the lowest 4-bit precision quantization
+        )
+        print("===> model loaded")
+
+        tokenizer = transformers.AutoTokenizer.from_pretrained('EleutherAI/gpt-neox-20b', padding_side="left", device_map="auto")
+            
+        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)
+
+        result = self.pipeline(inputs, **parameters)
+        print("===> result", result)
+
+        return result