utils.py

# -*- coding: utf-8 -*-
import torch
from zhipuai import ZhipuAI
import openai
import google.generativeai as genai
import requests
import json
import numpy as np
from bert_score import score
from transformers import AutoTokenizer
import re

def load_references(path):
    train_references = []
    valid_references = []
    test_references = []
    for name in ['train', 'valid', 'test']:
        file_name = name + '_per_suggestion.json'
        temp_references = []
        with open(path + file_name, 'r', encoding='utf-8') as f:
            datas = json.load(f)
            for data in datas:
                temp_references.append(data['output'])
        if name == 'train':
            train_references = temp_references
        elif name == 'valid':
            valid_references = temp_references
        else:
            test_references = temp_references
    return train_references, valid_references, test_references

def load_predictions(path):
    train_predictions = []
    valid_predictions = []
    test_predictions = []
    for name in ['train', 'valid', 'test']:
        file_name = name + '_prediction.jsonl'
        temp_predictions = []
        with open(path + file_name, 'r', encoding='utf-8') as f:
            for line in f:
                line_prediction = json.loads(line)['predict']
                start = line_prediction.find('[')
                end = line_prediction.find(']')
                prediction = line_prediction[start + 1:end]
                temp_predictions.append(prediction)

        if name == 'train':
            train_predictions = temp_predictions
        elif name =='valid':
            valid_predictions =temp_predictions
        else:
            test_predictions = temp_predictions
    return train_predictions, valid_predictions, test_predictions

def load_features(path, file_name):
    with open(path + file_name, 'r') as file:
        lines = file.readlines()
    all_data = []
    current_data = []
    for line in lines:
        parts = line.strip().split()
        if len(parts) == 3:
            if current_data:
                if len(current_data) < 45:
                    padding = [[0.0] * 35] * (45 - len(current_data))
                    current_data = padding + current_data
                all_data.append(current_data)
            current_data = []
        else:
            parts = line.strip().split(',')
            current_data.append([float(value) for value in parts])
    if current_data:
        if len(current_data) < 45:
            padding = [[0.0] * 35] * (45 - len(current_data))
            current_data = padding + current_data
        all_data.append(current_data)
    all_data_tensor = torch.tensor(all_data, dtype=torch.float32)
    return all_data_tensor

def load_labels(path, file_name):
    id_label_dict = {}

    with open(path + file_name, 'r') as file:
        for line in file:
            id_str, label_str = line.strip().split()
            id = int(id_str)
            label = int(label_str)
            id_label_dict[id] = label

    sorted_labels = [label for id, label in sorted(id_label_dict.items())]
    label_tensor = torch.tensor(sorted_labels, dtype=torch.int64)
    return label_tensor

def load_num2label(filename):
    result_dict = {}
    with open(filename, 'r') as file:
        for line in file:
            elements = line.strip().split()
            value, key = elements
            result_dict[int(key)] = value
    return result_dict

def get_similarity(target_feature, sim_feature):
    dot_product = torch.dot(target_feature, sim_feature)
    norm_v1 = torch.norm(target_feature)
    norm_v2 = torch.norm(sim_feature)
    similarity = dot_product / (norm_v1 * norm_v2)
    return similarity

def add_quadruple(quadruples, new_quadruple):
    min_fourth_value = min(quadruples, key=lambda x: x[1])
    min_index = quadruples.index(min_fourth_value)

    if new_quadruple[1] > min_fourth_value[1]:
        quadruples[min_index] = new_quadruple

def get_similar_patient(encoder_model, target_feature, train_features, train_labels, sim_num, device):
    ## label, similarity
    results = [(0, 0) for _ in range(sim_num)]

    for i in range(train_features.shape[0]):
        sim_feature = encoder_model(train_features[i].to(device))
        similarity = get_similarity(target_feature, sim_feature)
        temp_result = [train_labels[i], similarity]
        add_quadruple(results, temp_result)

    return results

def generate_prompt(prediction, sim_patients, id2label_dict, wo_self, wo_peer):
    if wo_peer:
        prompt = 'I am a patient. Please provide advice based on my prediction results and the data from patients similar to me.\n'
        disease_prompt = 'Based on my previous checkup results, the medical deep learning model predicts that my current condition is ' + prediction + '.\n'
        prompt = prompt + disease_prompt
    else:
        prompt = 'I am a patient. Please provide advice based on my prediction results and the data from patients similar to me.\n'
        disease_prompt = 'Based on my previous checkup results, the medical deep learning model predicts that my current condition is ' + prediction + '.\n'
        patient_prompt = 'In addition, the hospital analyzed the similarity of physical examination indicators and identified information about past patients with similar results to mine. Their discharge diagnoses are as follows (similarity range: 0-1): \n'
        for i in range(len(sim_patients)):
            disease = id2label_dict[sim_patients[i][0].item()]
            relevance = sim_patients[i][1].item()
            patient_prompt = patient_prompt + ('{}. The diagnosis result is {}, with a similarity score of {:.4f}.\n'.format(i + 1, disease, relevance))
        if not wo_self:
            prompt = prompt + disease_prompt + patient_prompt
        else:
            prompt = prompt + patient_prompt[3:]
    return prompt

def for_show_prompt(prediction, sim_patients, id2label_dict, wo_self, wo_peer):
    if wo_peer:
        prompt = 'I am a patient. Please provide advice based on my prediction results and the data from patients similar to me.\n'
        disease_prompt = 'Based on my previous checkup results, the medical deep learning model predicts that my current condition is ' + prediction + '.\n'
        prompt = prompt + disease_prompt
    else:
        prompt = 'I am a patient. Please provide advice based on my prediction results and the data from patients similar to me.\n'
        disease_prompt = 'Based on my previous checkup results, the medical deep learning model predicts that my current condition is ' + prediction + '.\n'
        patient_prompt = 'In addition, the hospital analyzed the similarity of physical examination indicators and identified information about past patients with similar results to mine. Their discharge diagnoses are as follows (similarity range: 0-1): \n'
        for i in range(len(sim_patients)):
            disease = id2label_dict[sim_patients[i][0].item()]
            relevance = sim_patients[i][1].item()
            patient_prompt = patient_prompt + ('{}. The diagnosis result is {}.\n'.format(i + 1, disease))
        if not wo_self:
            prompt = prompt + disease_prompt + patient_prompt
        else:
            prompt = prompt + patient_prompt[3:]
    return prompt

def generate_compare_prompt(data):
    data = data.cpu().numpy()
    data_str = '\n'.join(' '.join(map(str, data[i:i+35])) for i in range(0, len(data), 35))
    pattern = r'(^((0\.0\s){34}0\.0)\n)'
    while re.match(pattern, data_str):
        data_str = re.sub(pattern, '', data_str, count=1)
    prompt = 'I am a patient. Please provide advice based on my prediction results and the data from patients similar to me.\n' + \
             'The health examination indicators for patient include pH level, albumin, total protein, indirect bilirubin, direct bilirubin, total bilirubin, alkaline phosphatase, alanine aminotransferase, prealbumin, total bile acids, large platelet count, plateletcrit, large platelet ratio, mean platelet volume, platelet distribution width, red cell distribution width (coefficient of variation and standard deviation), absolute counts of basophils, eosinophils, monocytes, lymphocytes, and neutrophils, as well as their respective percentages. Additional indicators include mean corpuscular hemoglobin concentration, mean corpuscular hemoglobin, hematocrit, white blood cell count, platelet count, hemoglobin level, mean corpuscular volume, and red blood cell count.\n' \
             'Here are my health results from different dates: \n' + data_str
    return prompt

def modify_prompt(prompt, del_probs, offsets):
    del_index = np.random.choice(np.arange(len(del_probs)), size=1, p=del_probs.detach().cpu().numpy(), replace=False)[0]
    start, end = offsets[del_index]
    del_word = prompt[start:end]
    prompt = prompt[:start] + prompt[end:]
    return prompt, del_index, del_word

def truncate_text(text, max_length=512):
    tokenizer = AutoTokenizer.from_pretrained("./bert-base-uncased/")
    tokenized_text = tokenizer.tokenize(text)
    truncated_text = tokenizer.convert_tokens_to_string(tokenized_text[:max_length])
    return truncated_text

def get_reward(before_answer, after_answer, reference, device):
    # print([truncate_text(before_answer)])
    # print([truncate_text(reference)])
    _, _, before_score = score([before_answer], [reference], model_type="bert-base-uncased", lang="en", verbose=False, device=device)
    _, _, after_score = score([after_answer], [reference], model_type="bert-base-uncased", lang="en", verbose=False, device=device)
    # print('before bertscore:', before_score)
    # print('after bertscore:', after_score)
    reward = (after_score - before_score) * 100
    return reward, before_score, after_score

def loss_fcn(before_answer, after_answer, reference, del_probs, device):
    reward, _, _ = get_reward(before_answer, after_answer, reference, device)
    del_probs = del_probs.to(device)
    reward = reward.to(device)
    loss = -torch.sum(torch.log(del_probs)) * reward
    return loss

def get_llm_answer(prompt, llm):
    if 'glm' in llm:
        client = ZhipuAI(api_key="")
        response = client.chat.completions.create(
            model=llm,
            messages=[
                {"role": "user", "content": prompt},
            ],
        )
        return response.choices[0].message.content
    if 'gpt' in llm:
        openai.api_key = ""
        openai.api_base = "https://api.openai.com/v1"
        response = openai.ChatCompletion.create(
            model=llm,
            messages=[
                {"role": "system", "content": "You are an assistant with expertise in medical knowledge."},
                {"role": "user", "content": prompt}
            ]
        )
        return response['choices'][0]['message']['content']
    if 'gemini' in llm:
        genai.configure(api_key="")
        model = genai.get_model(llm)
        response = model.generate_message(
            messages=[
                {"role": "system", "content": "You are an assistant with expertise in medical knowledge."},
                {"role": "user", "content": prompt}
            ]
        )
        return response.last.message.content