Commit
·
6f045c2
1
Parent(s):
4225454
add custom handler
Browse files- handler.py +12 -12
handler.py
CHANGED
@@ -1,25 +1,24 @@
|
|
1 |
-
from typing import Dict,
|
2 |
-
from transformers import
|
3 |
|
4 |
class EndpointHandler:
|
5 |
def __init__(self, path="prashanthbsp/reasoning-cpg-entity-v1"):
|
6 |
-
#
|
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:
|
14 |
Return:
|
15 |
-
|
16 |
"""
|
17 |
-
# Extract
|
18 |
inputs = data.pop("inputs", data)
|
19 |
context = inputs.pop("context", inputs)
|
20 |
|
21 |
-
# Format prompt
|
22 |
-
prompt =
|
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 |
-
{
|
58 |
### Response:
|
59 |
-
<think>"""
|
60 |
|
61 |
-
#
|
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 |
}
|