atereoyinn
first commit
cd33ba8
raw
history blame
7.85 kB
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)