import re import os import simple_icd_10_cm as cm from transformers import AutoModelForCausalLM, AutoTokenizer # from openai import OpenAI from prompt_template import * from langchain_groq import ChatGroq from groq import Groq from dotenv import load_dotenv import csv import time load_dotenv() os.environ["LANGCHAIN_TRACING_V2"]="true" # groq_api_key=os.environ.get('GROQ_API_KEY') groq_api_key=os.getenv('GROQ_API_KEY') os.environ["LANGCHAIN_ENDPOINT"]="https://api.smith.langchain.com" LANGCHAIN_API_KEY=os.environ.get("LANGCHAIN_API_KEY") client = Groq() CHAPTER_LIST = cm.chapter_list def construct_translation_prompt(medical_note): """ Construct a prompt template for translating spanish medical notes to english. Args: medical_note (str): The medical case note. Returns: str: A structured template ready to be used as input for a language model. """ translation_prompt = """You are an expert Spanish-to-English translator. You are provided with a clinical note written in Spanish. You must translate the note into English. You must ensure that you properly translate the medical and technical terms from Spanish to English without any mistakes. Spanish Medical Note: {medical_note}""" return translation_prompt.format(medical_note = medical_note) def build_translation_prompt(input_note, system_prompt=""): """ Build a zero-shot prompt for translating spanish medical notes to english. Args: input_note (str): The input note or query. system_prompt (str): Optional initial system prompt or instruction. Returns: list of dict: A structured list of dictionaries defining the role and content of each message. """ input_prompt = construct_translation_prompt(input_note) return [{"role": "system", "content": system_prompt}, {"role": "user", "content": input_prompt}] def remove_extra_spaces(text): """ Remove extra spaces from a given text. Args: text (str): The original text string. Returns: str: The cleaned text with extra spaces removed. """ return re.sub(r'\s+', ' ', text).strip() def remove_last_parenthesis(text): """ Removes the last occurrence of content within parentheses from the provided text. Args: text (str): The input string from which to remove the last parentheses and its content. Returns: str: The modified string with the last parentheses content removed. """ pattern = r'\([^()]*\)(?!.*\([^()]*\))' cleaned_text = re.sub(pattern, '', text) return cleaned_text def format_code_descriptions(text, model_name): """ Format the ICD-10 code descriptions by removing content inside brackets and extra spaces. Args: text (str): The original text containing ICD-10 code descriptions. Returns: str: The cleaned text with content in brackets removed and extra spaces cleaned up. """ pattern = r'\([^()]*\)(?!.*\([^()]*\))' cleaned_text = remove_last_parenthesis(text) cleaned_text = remove_extra_spaces(cleaned_text) return cleaned_text def construct_prompt_template(case_note, code_descriptions, model_name): """ Construct a prompt template for evaluating ICD-10 code descriptions against a given case note. Args: case_note (str): The medical case note. code_descriptions (str): The ICD-10 code descriptions formatted as a single string. Returns: str: A structured template ready to be used as input for a language model. """ template = prompt_template_dict[model_name] return template.format(note=case_note, code_descriptions=code_descriptions) def build_zero_shot_prompt(input_note, descriptions, model_name, system_prompt=""): """ Build a zero-shot classification prompt with system and user roles for a language model. Args: input_note (str): The input note or query. descriptions (list of str): List of ICD-10 code descriptions. system_prompt (str): Optional initial system prompt or instruction. Returns: list of dict: A structured list of dictionaries defining the role and content of each message. """ if model_name == "llama3-70b-8192": code_descriptions = "\n".join(["* " + x for x in descriptions]) else: code_descriptions = "\n".join(["* " + x for x in descriptions]) input_prompt = construct_prompt_template(input_note, code_descriptions, model_name) return [{"role": "system", "content": system_prompt}, {"role": "user", "content": input_prompt}] def get_response(messages, model_name, temperature=0.0, max_tokens=500): """ Obtain responses from a specified model via the chat-completions API. Args: messages (list of dict): List of messages structured for API input. model_name (str): Identifier for the model to query. temperature (float): Controls randomness of response, where 0 is deterministic. max_tokens (int): Limit on the number of tokens in the response. Returns: str: The content of the response message from the model. """ response = client.chat.completions.create( model=model_name, messages=messages, temperature=temperature, max_tokens=max_tokens ) return response.choices[0].message.content def remove_noisy_prefix(text): # Removing numbers or letters followed by a dot and optional space at the beginning of the string cleaned_text = text.replace("* ", "").strip() cleaned_text = re.sub(r"^\s*\w+\.\s*", "", cleaned_text) return cleaned_text.strip() def parse_outputs(output, code_description_map, model_name): """ Parse model outputs to confirm ICD-10 codes based on a given description map. Args: output (str): The model output containing confirmations. code_description_map (dict): Mapping of descriptions to ICD-10 codes. Returns: list of dict: A list of confirmed codes and their descriptions. """ confirmed_codes = [] split_outputs = [x for x in output.split("\n") if x] for item in split_outputs: try: code_description, confirmation = item.split(":", 1) # print(confirmation) cnf,fact = confirmation.split(",", 1) if model_name == "llama3-70b-8192": code_description = remove_noisy_prefix(code_description) else: code_description = remove_noisy_prefix(code_description) if confirmation.lower().strip().startswith("yes"): try: code = code_description_map[code_description] confirmed_codes.append({"ICD Code": code, "Code Description": code_description,"Evidence From Notes":fact}) except Exception as e: # print(str(e) + " Here") continue except: continue return confirmed_codes def get_name_and_description(code, model_name): """ Retrieve the name and description of an ICD-10 code. Args: code (str): The ICD-10 code. Returns: tuple: A tuple containing the formatted description and the name of the code. """ full_data = cm.get_full_data(code).split("\n") return format_code_descriptions(full_data[3], model_name), full_data[1]