first commit

app.py

@@ -0,0 +1,208 @@
+import gradio as gr
+from faster_whisper import WhisperModel
+from pydantic import BaseModel, Field, AliasChoices, field_validator, ValidationError
+from typing import List
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+import csv
+import json
+import tempfile
+import torch
+
+
+# Initiate checkpoints for model loading
+numind_checkpoint = "numind/NuExtract-tiny"
+llama_checkpoint = "Atereoyin/Llama3_finetuned_for_medical_entity_extraction"
+whisper_checkpoint = "large-v3"
+
+quantization_config = BitsAndBytesConfig(
+    load_in_8bit=True,
+)
+
+# Load models with the correct device
+whisper_model = WhisperModel(whisper_checkpoint, device="cuda")
+numind_model = AutoModelForCausalLM.from_pretrained(numind_checkpoint, quantization_config=quantization_config, torch_dtype=torch.float16, trust_remote_code=True)
+numind_tokenizer = AutoTokenizer.from_pretrained(numind_checkpoint)
+llama_model = AutoModelForCausalLM.from_pretrained(llama_checkpoint, quantization_config=quantization_config, trust_remote_code=True)
+llama_tokenizer = AutoTokenizer.from_pretrained(llama_checkpoint)
+
+# Function to transcribe audio
+def transcribe_audio(audio_file_path):
+    try:
+        segments, info = whisper_model.transcribe(audio_file_path, beam_size=5)
+        text = "".join([segment.text for segment in segments])
+        return text
+    except Exception as e:
+        return str(e)
+
+# Functions for Person entity extraction
+def predict_NuExtract(model, tokenizer, text, schema, example=["","",""]):
+    schema = json.dumps(json.loads(schema), indent=4)
+    input_llm =  "<|input|>\n### Template:\n" +  schema + "\n"
+    for i in example:
+        if i != "":
+            input_llm += "### Example:\n"+ json.dumps(json.loads(i), indent=4)+"\n"
+
+    input_llm +=  "### Text:\n"+text +"\n<|output|>\n"
+    input_ids = tokenizer(input_llm, return_tensors="pt", truncation=True, max_length=4000).to("cuda")
+
+    output = tokenizer.decode(model.generate(**input_ids)[0], skip_special_tokens=True)
+    return output.split("<|output|>")[1].split("<|end-output|>")[0]
+
+
+#Function for generating promtps for Llama
+def prompt_format(text):
+    prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
+
+              ### Instruction:
+              {}
+
+              ### Input:
+              {}
+
+              ### Response:
+              {}"""
+
+    instruction = """Extract the following entities from the medical conversation:
+                    * **Symptoms:** List all the symptoms the patient mentions.
+                    * **Diagnosis:** List the doctor's diagnosis or potential diagnoses.
+                    * **Medical History:** Summarize the patient's relevant medical history.
+                    * **Action Plan:** List the recommended actions or treatment plan.
+
+                    Provide the result in the following JSON format:
+                    {
+                        "Symptoms": [...],
+                        "Diagnosis": [...],
+                        "Medical history": [...],
+                        "Action plan": [...]
+                    }"""
+    full_prompt = prompt.format(instruction, text, "")
+    return full_prompt
+
+
+#Pydantic Validator to validate Llama's response
+def validate_medical_record(response):
+
+    class MedicalRecord(BaseModel):
+        Symptoms: List[str] = Field(default_factory=list)
+        Diagnosis: List[str] = Field(default_factory=list)
+        Medical_history: List[str] = Field(
+            default_factory=list,
+            validation_alias=AliasChoices('Medical history', 'History of Patient')
+        )
+        Action_plan: List[str] = Field(
+            default_factory=list,
+            validation_alias=AliasChoices('Action plan', 'Plan of Action')
+        )
+
+        @field_validator('*', mode='before')
+        def ensure_list(cls, v):
+            if isinstance(v, str):
+                return [item.strip() for item in v.split(',')]
+            return v
+
+    try:
+        validated_data = MedicalRecord(**response)
+        return validated_data.dict()
+    except ValidationError as e:
+        return response
+
+
+
+# Function to predict medical entities using Llama
+def predict_Llama(model, tokenizer, text):
+    inputs = tokenizer(prompt_format(text), return_tensors="pt", truncation=True).to("cuda")
+
+    try:
+        outputs = model.generate(**inputs, max_new_tokens=128, temperature=0.2, use_cache=True)
+        extracted_entities = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+        response = extracted_entities.split("### Response:", 1)[-1].strip()
+        response_dict = {k.strip(): v.strip() for k, v in (line.split(': ', 1) for line in response.splitlines() if ': ' in line)}
+
+        validated_response = validate_medical_record(response_dict)
+
+        return validated_response
+    except Exception as e:
+        print(f"Error during Llama prediction: {str(e)}")
+        return {}
+
+
+#Control function that cordinates communication of other functions to map entities to form fields
+def process_audio(audio):
+    if isinstance(audio, str):
+        with open(audio, 'rb') as f:
+            audio_bytes = f.read()
+    else:
+        audio_bytes = audio
+
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_audio:
+        temp_audio.write(audio_bytes)
+        temp_audio.flush()
+        audio_path = temp_audio.name
+
+    transcription = transcribe_audio(audio_path)
+
+    person_schema = """{"Name": "","Age": "","Gender": ""}"""
+    person_entities_raw = predict_NuExtract(numind_model, numind_tokenizer, transcription, person_schema)
+
+    try:
+        person_entities = json.loads(person_entities_raw)
+    except json.JSONDecodeError as e:
+        return f"Error in NuExtract response: {str(e)}"
+
+    medical_entities = predict_Llama(llama_model, llama_tokenizer, transcription)
+
+    return (
+        person_entities.get("Name", ""),
+        person_entities.get("Age", ""),
+        person_entities.get("Gender", ""),
+        ", ".join(medical_entities.get("Symptoms", [])),
+        ", ".join(medical_entities.get("Diagnosis", [])),
+        ", ".join(medical_entities.get("Medical_history", [])),
+        ", ".join(medical_entities.get("Action_plan", []))
+    )
+
+
+
+#Function that allows users to download information
+def download_csv(name, age, gender, symptoms, diagnosis, medical_history, action_plan):
+    csv_file = tempfile.NamedTemporaryFile(delete=False, suffix=".csv")
+  
+    with open(csv_file.name, mode='w', newline='') as file:
+        writer = csv.writer(file)
+        writer.writerow(["Name", "Age", "Gender", "Symptoms", "Diagnosis", "Medical History", "Plan of Action"])
+        writer.writerow([name, age, gender, symptoms, diagnosis, medical_history, action_plan])
+    
+    return csv_file.name
+
+
+
+# Gradio interface to create a web-based form for users to input audio and fill the medical diagnostic form
+demo = gr.Interface(
+    fn=process_audio,
+    inputs=[
+        gr.Audio(type="filepath")
+    ],
+    outputs=[
+        gr.Textbox(label="Name"),
+        gr.Textbox(label="Age"),
+        gr.Textbox(label="Gender"),
+        gr.Textbox(label="Symptoms"),
+        gr.Textbox(label="Diagnosis"),
+        gr.Textbox(label="Medical History"),
+        gr.Textbox(label="Plan of Action"),
+    ],
+    title="Medical Diagnostic Form Assistant",
+    description="Upload an audio file or record audio to generate a medical diagnostic form."
+)
+
+with demo:
+    download_button = gr.Button("Download CSV")
+    download_button.click(
+        fn=lambda name, age, gender, symptoms, diagnosis, medical_history, action_plan: download_csv(name, age, gender, symptoms, diagnosis, medical_history, action_plan),
+        inputs=demo.output_components,
+        outputs=gr.File(label="Download CSV")
+    )
+
+demo.launch(share=True)
+

requirements.txt

@@ -0,0 +1,5 @@
+gradio
+transformers
+torch
+faster-whisper
+pydantic