File size: 4,536 Bytes
641eed3 c72b4ed 641eed3 595788e 641eed3 1fe53f2 641eed3 1fe53f2 641eed3 1fe53f2 641eed3 1fe53f2 abf9374 641eed3 abf9374 641eed3 abf9374 641eed3 | 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 97 98 99 100 101 102 103 104 105 106 107 108 109 | from typing import Any, Dict, List
import torch
import json
import transformers
import bitsandbytes
from transformers import AutoModelForCausalLM, AutoTokenizer
dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] == 8 else torch.float16
class EndpointHandler:
def __init__(self, path=""):
print("ANSHAY calling EndpointHandler __init__")
self.tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
self.model = AutoModelForCausalLM.from_pretrained(
path,
return_dict=True,
device_map="auto",
load_in_8bit=True,
torch_dtype=dtype,
trust_remote_code=True,
)
print("ANSHAY got model and tokenizer")
print(self.model)
print(self.tokenizer)
generation_config = self.model.generation_config
generation_config.max_new_tokens = 3
generation_config.temperature = 0.1
generation_config.pad_token_id = self.tokenizer.eos_token_id
generation_config.eos_token_id = self.tokenizer.eos_token_id
self.generation_config = generation_config
self.pipeline = transformers.pipeline(
"text-generation", model=self.model, tokenizer=self.tokenizer
)
print("ANSHAY got pipeline")
print(self.pipeline)
def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
print("ANSHAY in __call__ of handler")
print("ANSHAY printing data:")
print(data)
input_string = data.pop("inputs");
print("ANSHAY printing input_string:")
print(input_string)
input_dict = json.loads(input_string)
print("ANSHAY printing input_dict:")
print(input_dict)
gender = input_dict["gender"]
age = input_dict["age"]
occupation = input_dict["occupation"]
sleep_hours = input_dict["sleep_hours"]
physical_activity = input_dict["physical_activity"]
caffiene_intake = input_dict["caffiene_intake"]
alcohol_consumption = input_dict["alcohol_consumption"]
smoking = input_dict["smoking"]
family_anxiety = input_dict["family_anxiety"]
stress_level = input_dict["stress_level"]
heart_rate = input_dict["heart_rate"]
breathing_rate = input_dict["breathing_rate"]
sweating_level = input_dict["sweating_level"]
dizziness = input_dict["dizziness"]
medication = input_dict["medication"]
therapy_sessions = input_dict["therapy_sessions"]
recent_life_event = input_dict["recent_life_event"]
diet_quality = input_dict["diet_quality"]
chat_formatted = [{
"role": "system",
"content": "You are an anxiety calculating bot. You help people rate their anxiety on a scall of 1-10, after understanding their current lifestyle.",
},
{
"role": "user",
"content": f"Here are my details:\nGender:{gender}\nAge:{age}\nOccupation:{occupation}\nSleep Hours:{sleep_hours}\nPhysical Activity (hrs/week):{physical_activity}\nCaffeine Intake (mg/day):{caffiene_intake}\nAlcohol Consumption (drinks/week):{alcohol_consumption}\nSmoking:{smoking}\nFamily History of Anxiety:{family_anxiety}\nStress Level (1-10):{stress_level}\nHeart Rate (bpm):{heart_rate}\nBreathing Rate (breaths/min):{breathing_rate}\nSweating Level (1-5):{sweating_level}\nDizziness:{dizziness}\nMedication:{medication}\nTherapy Sessions (per month):{therapy_sessions}\nRecent Major Life Event:{recent_life_event}\nDiet Quality (1-10):{diet_quality}"
},
{
"role": "assistant",
"content": f"Anxiety Rating:"
}]
print("ANSHAY printing final chat formatted:")
print(chat_formatted)
chat_template_applied = self.tokenizer.apply_chat_template(chat_formatted, tokenize=False, continue_final_message=True)
print("ANSHAY printing final chat_template_applied:")
print(chat_template_applied)
result = self.pipeline(chat_template_applied, generation_config=self.generation_config)
print("ANSHAY printing final result:")
print(result)
generated_text = result[0]['generated_text']
print("ANSHAY printing generated_text:")
print(generated_text)
anxiety_rating = generated_text[-3:]
print("ANSHAY printing anxiety_rating:")
print(anxiety_rating)
return anxiety_rating |