lifeofcoding commited on
Commit
619846d
·
verified ·
1 Parent(s): 7ae55b3

Delete handler.py

Browse files
Files changed (1) hide show
  1. 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