Delete handler.py
Browse files- handler.py +0 -45
handler.py
DELETED
@@ -1,45 +0,0 @@
|
|
1 |
-
from typing import Any, Dict, List
|
2 |
-
|
3 |
-
import torch
|
4 |
-
import transformers
|
5 |
-
from transformers import AutoModelForCausalLM, AutoTokenizer
|
6 |
-
|
7 |
-
dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] == 8 else torch.float16
|
8 |
-
|
9 |
-
class EndpointHandler:
|
10 |
-
def __init__(self, path=""):
|
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=False,
|
17 |
-
torch_dtype=dtype,
|
18 |
-
trust_remote_code=True,
|
19 |
-
)
|
20 |
-
|
21 |
-
generation_config = model.generation_config
|
22 |
-
generation_config.max_new_tokens = 10000
|
23 |
-
generation_config.temperature = 0.01
|
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", model=model, tokenizer=tokenizer
|
31 |
-
)
|
32 |
-
|
33 |
-
def __call__(self, data: Dist[str, Any]) -> Dict[str, Any]:
|
34 |
-
prompt = data.pop("inputs", data)
|
35 |
-
|
36 |
-
system_message = """
|
37 |
-
You are an expert at analyzing the text of narrative scenes. Only respond with complete JSON responses beginning with [ and end with ]. Do not add new lines Separate the text of the scene to analyze into passages spoken aloud by a character from text not spoken aloud by a character. Identify the character speaking, and the tone being used. To help identify dialog, we have turned all quotation marks into the string SPOOOOOOKEN. For example the sentence SPOOOOOOKENHello, Janet,SPOOOOOOKEN John growled angrily. would result in the following output format: [{"position": 1, "type": "dialog", "narrator": "John", "tone": "angry", "content": "SPOOOOOOKENHello, Janet,SPOOOOOOKEN "}, {"position": 2, "type": "narration", "narrator": "Narrator", "tone": "", "content": "John growled angrily."}] If there is a new-line character, separate it as its own json object. Here is the scene to analyze:
|
38 |
-
""".strip()
|
39 |
-
|
40 |
-
messages = [
|
41 |
-
{"role": "system", "content": system_message},
|
42 |
-
{"role": "user", "content": prompt},
|
43 |
-
]
|
44 |
-
result = self.pipeline(messages, generation_config=self.generation_config)
|
45 |
-
return result
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|