Spaces:
Running on Zero
Running on Zero
| import spaces | |
| import gradio as gr | |
| import torch | |
| import torch.nn.functional as F | |
| from transformers import DistilBertTokenizerFast, DistilBertForSequenceClassification | |
| # This makes your code work on both your local GPU and the HF CPU Space | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| # 1. Load Model & Tokenizer | |
| model_path = "./assets/model" | |
| tokenizer = DistilBertTokenizerFast.from_pretrained(model_path) | |
| model = DistilBertForSequenceClassification.from_pretrained(model_path) | |
| model.to(device) | |
| # Set model to evaluation mode | |
| model.eval() | |
| # Mapping for dropdown options | |
| SAMPLE_DATA = { | |
| "Normal 1": {"text": "Well, there's a mother who is standing in the kitchen. She is washing dishes at the sink, but the sink is overflowing and water is spilling onto the floor. Meanwhile, the two children, a boy and a girl, are reaching for a cookie jar on a high shelf. The boy is standing on a stool that looks like it's about to tip over.", | |
| "label": "Cognitively Normal"}, | |
| "Normal 2": {"text": "It's a scene in a kitchen where the woman is doing the dishes, and she doesn't seem to notice that the water is running over the basin. The kids are trying to get cookies out of the jar. The boy has grabbed the cookies and is passing one to his sister. However, he is reaching out for more and while he has another cookie in his hand, he is about to fall down from the stool.the little girl is reaching up, looking like she's asking for a cookie too.", "label": "Cognitively Normal"}, | |
| "Impaired 1": {"text": "Uh, dishes... water... floor. Mom. Cookie... jar. Boy, girl. Stool falling. Yeah, cookie.", | |
| "label": "Impaired"}, | |
| "Impaired 2": {"text": "The... the thing is, the water is... it is wet. The children, they want the, the sweets. One is up, up on the, the chair. It's falling. Not good.", | |
| "label": "Impaired"} | |
| } | |
| # 2. Clinical Class Names | |
| CLASS_NAMES = ["Cognitively Normal", "Impaired"] | |
| # 3. Prediction Function | |
| def process_input(input_mode, dropdown_choice, file_obj): | |
| # Logic: Prioritize based on the selected mode | |
| if input_mode == "Select from Dropdown" and dropdown_choice: | |
| text = SAMPLE_DATA[dropdown_choice]["text"] | |
| true_label = SAMPLE_DATA[dropdown_choice]["label"] | |
| elif input_mode == "Upload File" and file_obj: | |
| # file_obj is a TemporaryFile path object in Gradio | |
| with open(file_obj.name, 'r', encoding='utf-8') as f: | |
| text = f.read() | |
| true_label = "N/A (Uploaded File)" | |
| else: | |
| return "Please select or upload a valid file.", {}, "No text provided", "N/A" | |
| # Inference | |
| inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512).to(device) | |
| with torch.no_grad(): | |
| logits = model(**inputs).logits | |
| probs = F.softmax(logits, dim=1).squeeze().tolist() | |
| prob_dict = {"Cognitively Normal": probs[0], "Impaired": probs[1]} | |
| pred_class = "Cognitively Normal" if probs[0] > probs[1] else "Impaired" | |
| return pred_class, prob_dict, text, true_label | |
| # 4. Custom CSS for width limits | |
| custom_css = """ | |
| #narrow-container { | |
| max-width: 1200px; | |
| margin: 0 auto; | |
| } | |
| max-width:1200px | |
| """ | |
| # 5. Build the UI | |
| with gr.Blocks() as demo: | |
| # Wrap the header in a column with the custom CSS ID | |
| with gr.Column(elem_id="narrow-container"): | |
| gr.Markdown("# Boston CCTD Task: Cognitive Impairment Classification") | |
| # Task Description | |
| gr.Markdown( | |
| "This model classifies patient transcripts based on the **Boston Cookie Theft Picture Description Task (CCTD)**. " | |
| "In this assessment, subjects are asked to describe the events occurring in the standardized image below. " | |
| "The model analyzes the linguistic patterns, vocabulary, and syntax within their spoken transcripts to differentiate " | |
| "between **Cognitively Normal** and **Impaired** individuals." | |
| ) | |
| gr.Markdown( | |
| """ | |
| ### Project Details & Methodology | |
| * **Model Architecture**: Utilized `DistilBertTokenizerFast` and `DistilBertForSequenceClassification`. | |
| * **Training Source**: Training code was adopted from [PROCESS-2/codes/PROCESS2_BASELINE_LLM.py](https://github.com/Madhurananda/PROCESS-2). | |
| * **Dataset**: Powered by the [CognoSpeak/PROCESS-2](https://huggingface.co/datasets/CognoSpeak/PROCESS-2) dataset. | |
| * **Task Specifics**: Implementation is limited strictly to the Cookie Theft Task (CTD). | |
| * **Classification**: Performs 2-class classification, merging MCI (Mild Cognitive Impairment) and Dementia categories into a single 'Cognitively Impaired' group. | |
| * **Input Data**: Utilizes manual transcripts for classification accuracy. | |
| """ | |
| ) | |
| # Image Component - PASTE YOUR PROJECT LINK HERE | |
| project_image_url = "./assets/cookie_theft.png" | |
| gr.Image( | |
| value=project_image_url, | |
| interactive=False, | |
| label="Standardized Assessment Image" | |
| ) | |
| # Visual divider | |
| gr.HTML("<hr style='margin-top: 20px; margin-bottom: 20px;' />") | |
| with gr.Row(): | |
| # Left side: File upload | |
| with gr.Column(): | |
| input_mode = gr.Radio(["Upload File", "Select from Dropdown"], label="Choose Input Method", value="Upload File") | |
| with gr.Column(): | |
| file_input = gr.File(label="Upload Transcript") | |
| dropdown_input = gr.Dropdown(list(SAMPLE_DATA.keys()), label="Select Sample", visible=False) | |
| # Toggle visibility | |
| input_mode.change(lambda x: (gr.update(visible=x=="Upload File"), gr.update(visible=x=="Select from Dropdown")), | |
| inputs=input_mode, outputs=[file_input, dropdown_input]) | |
| submit_btn = gr.Button("Analyze Transcript", variant="primary") | |
| # with gr.Column(): | |
| # transcript_output = gr.Textbox( | |
| # label="Extracted Transcript Text", | |
| # lines=10, # Sets a default visible height | |
| # max_lines=15, # Caps the height and enables scrolling if text is longer | |
| # interactive=False # Prevents the user from typing in it | |
| # ) | |
| text_out = gr.Textbox(label="Transcript Content", lines=5) | |
| with gr.Row(): | |
| pred_out = gr.Textbox(label="Predicted Diagnosis") | |
| truth_out = gr.Textbox(label="Ground Truth Label") | |
| prob_out = gr.Label(label="Probabilities") | |
| submit_btn.click( | |
| process_input, | |
| inputs=[input_mode, dropdown_input, file_input], | |
| outputs=[pred_out, prob_out, text_out, truth_out] | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(theme=gr.themes.Cyberpunk(), css=custom_css) |