muzammil-eds commited on
Commit
4c9c2cf
1 Parent(s): 745f903

Create handler.py

Browse files
Files changed (1) hide show
  1. handler.py +38 -0
handler.py ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Dict, List, Any
2
+ from transformers import AutoTokenizer, AutoModelForCausalLM
3
+ import transformers
4
+ import torch
5
+
6
+ dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] == 8 else torch.float16
7
+
8
+ class EndpointHandler:
9
+ def __init__(self, path=""):
10
+
11
+ tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
12
+ model = AutoModelForCausalLM.from_pretrained(
13
+ path,
14
+ return_dict=True,
15
+ device_map="auto",
16
+ load_in_8bit=True,
17
+ torch_dtype=dtype,
18
+ trust_remote_code=True,
19
+ )
20
+
21
+ generation_config = model.generation_config
22
+ generation_config.max_new_tokens = 60
23
+ generation_config.temperature = 0
24
+ generation_config.num_return_sequences = 1
25
+ generation_config.pad_token_id = tokenizer.eos_token_id
26
+ generation_config.eos_token_id = tokenizer.eos_token_id
27
+ self.generation_config = generation_config
28
+
29
+ self.pipeline = transformers.pipeline(
30
+ "text-generation",
31
+ model=model,
32
+ tokenizer=tokenizer
33
+ )
34
+
35
+ def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
36
+ prompt = data.pop("inputs", data)
37
+ result = self.pipeline(prompt, generation_config=self.generation_config)
38
+ return result