mkeohane01 commited on
Commit
16917d3
·
verified ·
1 Parent(s): 6adf070

Created handler.py

Browse files
Files changed (1) hide show
  1. handler.py +50 -0
handler.py ADDED
@@ -0,0 +1,50 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Any, Dict
2
+ import torch
3
+ from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, pipeline
4
+
5
+ dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] == 8 else torch.float16
6
+
7
+ class EndpointHandler:
8
+ def __init__(self, path=""):
9
+
10
+ quantization_config = BitsAndBytesConfig(load_in_8bit=True)
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
+ torch_dtype = dtype,
17
+ trust_remote_code = True,
18
+ quantization_config=quantization_config
19
+ )
20
+
21
+ gen_config = model.generation_config
22
+ gen_config.max_new_tokens = 256
23
+ gen_config.num_return_sequences = 1
24
+ gen_config.pad_token_id = tokenizer.eos_token_id
25
+ gen_config.eos_token_id = tokenizer.eos_token_id
26
+
27
+ self.generation_config = gen_config
28
+
29
+ self.pipeline = pipeline(
30
+ 'text-generation', model=model, tokenizer=tokenizer
31
+ )
32
+
33
+
34
+
35
+ def __call__(self, data: Dict[dict, Any]) -> Dict[str, Any]:
36
+ prompt = data.pop("inputs", data)
37
+
38
+ instruction = "Create a list of chords,a corresponding scale to improve with, title, and style along with an example in ABC notation based on this input in JSON format."
39
+
40
+ full_prompt = f"""<s>
41
+ ### Instruction:
42
+ {instruction}
43
+ ### Input:
44
+ {prompt}
45
+ ### Response:
46
+ """
47
+
48
+ result = self.pipeline(full_prompt, generation_config = self.generation_config)[0]
49
+
50
+ return result