usamakenway commited on
Commit
ecaa798
1 Parent(s): 06e5302

add custom pipeline

Browse files
Files changed (1) hide show
  1. pipleline.py +29 -0
pipleline.py ADDED
@@ -0,0 +1,29 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Dict, List, Any
2
+ from transformers import AutoTokenizer, TextGenerationPipeline, pipeline
3
+ from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
4
+
5
+
6
+ class PreTrainedPipeline():
7
+ def __init__(self, path=""):
8
+ tokenizer = AutoTokenizer.from_pretrained(path)
9
+ model = AutoGPTQForCausalLM.from_quantized(path, device="cuda:0", use_safetensors=True)
10
+
11
+
12
+ def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
13
+ """
14
+ Args:
15
+ data (:obj:):
16
+ includes the input data and the parameters for the inference.
17
+ Return:
18
+ A :obj:`list`:. The list contains the embeddings of the inference inputs
19
+ """
20
+ inputs = data.get("inputs", data)
21
+ parameters = data.get("parameters", {})
22
+
23
+ # tokenize the input
24
+ input_ids = self.tokenizer(inputs,return_tensors="pt").input_ids.to(self.model.device)
25
+ # run the model
26
+ logits = self.model.generate(input_ids, **parameters)
27
+ # Perform pooling
28
+ # postprocess the prediction
29
+ return {"generated_text": self.tokenizer.decode(logits[0].tolist())}