model.py

import torch
from peft import LoraConfig, AutoPeftModelForCausalLM, prepare_model_for_kbit_training, get_peft_model
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, GPTQConfig, TrainingArguments
#from trl import SFTTrainer
import os
import json
import pandas as pd


{
    'Text': "Ms. Powelton is a 20-year-old female who presents to the ED with abdominal pain. The pain began 8-10 hours ago, originated in the RLQ, woke her from sleep, and is described as dull, achy, a bad cramp, and constant (5/10, getting worse). She denies radiation of pain, denies pain elsewhere on the body. Ibuprofen is not helping as usual. She has a history of 34 episodes like this over the past 6 months, but nothing as bad as this episode. Previous episodes resolved on their own. She denies nausea/vomiting. LMP was 2 weeks ago. Regular sexually active 9 months ago, uses condoms. Decreased appetite, hasn't had food since last night but still drinking water. Diarrhea for the past 3 days. Denies fevers/chills. No chest pain, SOB, cough, palpitations, headache, dizziness, changes in urination. Denies any drug or cigarette use. Drinks alcohol occasionally. Parents are both in good health. No family history. No surgical or medical history.",
    'Question': "What specific nursing actions would you take based on the patient's presentation?",
    
    'Abdominal Pain': 
        { "right": "['1) Perform abdominal palpation to assess tenderness and check for rebound tenderness', '2) Initiate monitoring of vital signs to observe for any signs of instability', ....]",
    	  "wrong": "['1) Administer pain medication immediately without further assessment', '2) Order a chest X-ray to rule out respiratory causes of pain', ....]" },
    
    'Diarrhea for the Past 3 Days' : 
        { 'right': "['1) Inquire about the character, frequency, and color of the diarrhea to assess for potential causes', '2) Obtain a stool sample for laboratory analysis to identify infectious etiologies', ....]",
          'wrong': "['1) Suggest a high-fiber diet to alleviate diarrhea symptoms', '2) Refer the patient for an immediate colonoscopy to investigate the diarrhea']" },
        
    'History of Similar Episodes' : 
        { 'right': "['1) Explore the patient's history of similar episodes, including triggers and outcomes', '2) Order imaging studies such as an abdominal ultrasound to investigate recurrent abdominal pain', ....]",
           'wrong': "['1) Advise the patient to ignore the recurring episodes as they seem self-resolving', '2) Prescribe a long-term pain management plan without further investigation', ....]" },
    
    'Recent Sexual Activity and Condom Use' : 
        { 'right': "['1) Discuss sexual history to assess for potential sexually transmitted infections', '2) Provide education on contraceptive methods and safe sexual practices', ....]",
          'wrong': "['1) Disregard sexual history as it is irrelevant to the abdominal pain', '2) Immediately order a comprehensive STI panel without patient discussion', ....]" },
}




compute_dtype = getattr(torch, "float16")

model_path="/home/ubuntu/item_multiple_group/llama-2-7b.Q4_K_M.gguf",

quant_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=compute_dtype,
        bnb_4bit_use_double_quant=False,
    )


model = AutoModelForCausalLM.from_pretrained(
        model_path, 
        quantization_config=quant_config,
        device_map={"": 0}
    )

def generate_text_from_prompt(prompt,
                             max_tokens=200,
                             temperature=0.3,
                             top_p=0.1,
                             echo=True,
                             stop=["Q", "\n"]):

    # Define the parameters
    model_output = model(
        prompt,
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
        echo=echo,
        stop=stop,
    )

    return model_output


# Example usage:
custom_dict = {
    'Question': '.........?',
    'Key-1': {'right': '...', 'wrong': '...'},
    'Key-2': {'right': '...', 'wrong': '...'},
    'Key-3': {'right': '...', 'wrong': '...'},
    # Add other keys and values as needed
}

def create_prompt(custom_dict):
    """Generate a prompt with system message and user message based on custom_dict."""
    
    system_message = "You are a helpful, respectful, and honest subject matter expert in medical and \n" \
                     "health care domain. For the given patient's medical {text}, you need to generate contextual " \
                     f"based keywords Always answer as helpfully as possible, " \
                     "while being safe. Please ensure that your responses are socially unbiased and positive " \
                     "in nature.\n"

    options = ""
        
    for key, values in custom_dict.items():
        options += f"    '{key}': {{\n"
        options += f"        '{key}(right)': {values['right']},\n"
        options += f"        '{key}(wrong)': {values['wrong']},\n"
        options += f"    }},\n"


    formatted_prompt = f"<s>[INST] <<SYS>>\n{system_message}\n<</SYS>> please provide output as per the below optiions formate\n\nOutput:{options} [/INST]"
    return formatted_prompt


prompt = create_prompt(custom_dict)
print(prompt)


# Generate text from the modified prompt
model_response = generate_text_from_prompt(prompt)

print(model_response)