Update helpers.py

helpers.py

@@ -1,214 +1,215 @@
-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')
-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]
-
+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]
+