File size: 3,716 Bytes
d847823 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 | import torch
import re
from typing import Any, Dict
from transformers import AutoModelForCausalLM, AutoTokenizer
class EndpointHandler:
def __init__(self, path=""):
# load model and tokenizer from path
self.tokenizer = AutoTokenizer.from_pretrained(path)
self.model = AutoModelForCausalLM.from_pretrained(
path, device_map="auto", torch_dtype=torch.float16, trust_remote_code=True
)
self.device = "cuda" if torch.cuda.is_available() else "cpu"
def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
# process input
inputs = data.pop("inputs", data)
parameters = data.pop("parameters", None)
prompt = f'''Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n'''
prompt += f'### Instruction:\n'
prompt += f'''Three telecommunications experts analyze the behavior of a phone number based on its call records to evaluate whether it is suspicious behavior. The three experts will discuss together and provide the most confident probability value (0%–100%) of "Yes". No explanation is required.\n\n'''
prompt += f"### Question:\n{inputs}"
prompt += f"### Response:\n"
self.tokenizer.pad_token = self.tokenizer.eos_token
generation_config = {
'max_new_tokens':128,
'top_p':None,
'do_sample':False,
'num_beams': 5,
'temperature':None
}
self.model.generation_config.pad_token_id = self.tokenizer.pad_token_id
llama3_prompt = [
{
"role": "system",
"content": "Below is an instruction that describes a task. Write a response that appropriately completes the request."
},
{
"role": "user",
"content": ""
}
]
# preprocess
# inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
llama3_prompt[1]['content'] = prompt
input_ids = self.tokenizer.apply_chat_template(llama3_prompt, add_generation_prompt=True, add_special_tokens=True, return_tensors="pt").to(self.device)
prompt_len = input_ids.shape[-1]
outputs = self.model.generate(input_ids, **generation_config)
generated_answer = self.tokenizer.decode(outputs[0, prompt_len:], skip_special_tokens=True)
p = self.getProbability(generated_answer)
if p > 50.0:
generated_answer = 'Yes'
else:
generated_answer = 'No'
prediction =[{"generated_text": generated_answer}]
print("--prediction--",prediction )
return prediction
def getProbability(text):
match = re.search(r'(\d+(?:\.\d+)?)%', text)
if match:
number = float(match.group(1))
else:
# 如果沒有 '%'，提取第一個數字
match = re.search(r'(\d+(?:\.\d+)?)', text)
if match:
number = float(match.group(1))
else:
number = None # 未找到數字
if number is not None:
print(number)
return number
else:
# 如果沒有數字，找到第一個 'yes' 或 'no'，不區分大小寫
match = re.search(r'\b(yes|no)\b', text, re.IGNORECASE)
if match:
# 輸出匹配到的 'yes' 或 'no'，保持原始大小寫
print(match.group(0))
if match.group(0).lower() == 'yes':
return 100.0
else:
return 0.0
else:
return 0.0 |