prashanthbsp commited on
Commit
6f045c2
·
1 Parent(s): 4225454

add custom handler

Browse files
Files changed (1) hide show
  1. handler.py +12 -12
handler.py CHANGED
@@ -1,25 +1,24 @@
1
- from typing import Dict, List, Any
2
- from transformers import AutoModelForCausalLM, AutoTokenizer
3
 
4
  class EndpointHandler:
5
  def __init__(self, path="prashanthbsp/reasoning-cpg-entity-v1"):
6
- # Standard HF model loading - compatible with TGI
7
  self.tokenizer = AutoTokenizer.from_pretrained(path)
8
- # Model is loaded by the TGI server, not by the handler
9
 
10
  def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
11
  """
12
  data args:
13
- inputs: text or dict containing text
14
  Return:
15
- A dict with the model's response
16
  """
17
- # Extract inputs
18
  inputs = data.pop("inputs", data)
19
  context = inputs.pop("context", inputs)
20
 
21
- # Format prompt according to your requirements
22
- prompt = f"""Below is an instruction that describes a task, paired with an input that provides further context.
23
  Write a response that appropriately completes the request.
24
  Before answering, think carefully about the task to ensure a logical and accurate response.
25
 
@@ -54,15 +53,16 @@ class EndpointHandler:
54
  }}
55
 
56
  ### Social Media Post:
57
- {context}
58
  ### Response:
59
- <think>"""
60
 
61
- # For TGI, we return a dict with the prompt and generation params
62
  return {
63
  "inputs": prompt,
64
  "parameters": {
65
  "max_new_tokens": 1200,
 
66
  "do_sample": False,
67
  "return_full_text": False # Only return the generated text, not the prompt
68
  }
 
1
+ from typing import Dict, Any
2
+ from transformers import AutoTokenizer
3
 
4
  class EndpointHandler:
5
  def __init__(self, path="prashanthbsp/reasoning-cpg-entity-v1"):
6
+ # Only load the tokenizer - the model is loaded by TGI
7
  self.tokenizer = AutoTokenizer.from_pretrained(path)
 
8
 
9
  def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
10
  """
11
  data args:
12
+ inputs: Text or dict containing text
13
  Return:
14
+ Dict with prompt and generation parameters
15
  """
16
+ # Extract the input text
17
  inputs = data.pop("inputs", data)
18
  context = inputs.pop("context", inputs)
19
 
20
+ # Format the prompt
21
+ prompt = """Below is an instruction that describes a task, paired with an input that provides further context.
22
  Write a response that appropriately completes the request.
23
  Before answering, think carefully about the task to ensure a logical and accurate response.
24
 
 
53
  }}
54
 
55
  ### Social Media Post:
56
+ {0}
57
  ### Response:
58
+ <think>""".format(context)
59
 
60
+ # Return the formatted prompt and generation parameters for TGI
61
  return {
62
  "inputs": prompt,
63
  "parameters": {
64
  "max_new_tokens": 1200,
65
+ "temperature": 0.01, # Low temperature for more deterministic outputs
66
  "do_sample": False,
67
  "return_full_text": False # Only return the generated text, not the prompt
68
  }