| 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 |
