Britania commited on
Commit
ae146ef
·
verified ·
1 Parent(s): 17520ba

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +10 -2
handler.py CHANGED
@@ -4,12 +4,20 @@ import torch
4
  import json
5
 
6
  def model_inference(model, tokenizer, text, **kwargs):
 
 
 
 
 
 
 
 
7
  inputs = tokenizer(text, return_tensors="pt").to("cuda:0")
 
8
  with torch.inference_mode():
9
- outputs = model.generate(**inputs, **kwargs)
10
  return tokenizer.decode(outputs[0][inputs_length:], skip_special_tokens=True)
11
 
12
-
13
  class EndpointHandler():
14
  def __init__(self, path=""):
15
  self.model = AutoModelForCausalLM.from_pretrained(
 
4
  import json
5
 
6
  def model_inference(model, tokenizer, text, **kwargs):
7
+ default_kwargs = {
8
+ "max_new_tokens": 1024,
9
+ "temperature": 0.5,
10
+ "top_p": 0.95,
11
+ "top_k": 150,
12
+ "repetition_penalty": 1.4
13
+ }
14
+ default_kwargs.update(kwargs)
15
  inputs = tokenizer(text, return_tensors="pt").to("cuda:0")
16
+ inputs_length = len(inputs["input_ids"][0])
17
  with torch.inference_mode():
18
+ outputs = model.generate(**inputs, **default_kwargs)
19
  return tokenizer.decode(outputs[0][inputs_length:], skip_special_tokens=True)
20
 
 
21
  class EndpointHandler():
22
  def __init__(self, path=""):
23
  self.model = AutoModelForCausalLM.from_pretrained(