Spaces:
Running
Running
| import numpy as np | |
| import torch | |
| from transformers import AutoTokenizer, AutoModelForCausalLM | |
| import gradio as gr | |
| # Model setup | |
| device = torch.device('cpu') # Use 'cuda' if GPU is available | |
| dtype = torch.float32 # Data type for model processing | |
| model_name_or_path = 'GoodBaiBai88/M3D-LaMed-Phi-3-4B' | |
| proj_out_num = 256 | |
| # Load model and tokenizer | |
| model = AutoModelForCausalLM.from_pretrained( | |
| model_name_or_path, | |
| torch_dtype=torch.float32, | |
| device_map='cpu', | |
| trust_remote_code=True | |
| ) | |
| tokenizer = AutoTokenizer.from_pretrained( | |
| model_name_or_path, | |
| model_max_length=512, | |
| padding_side="right", | |
| use_fast=False, | |
| trust_remote_code=True | |
| ) | |
| # Chat history storage | |
| chat_history = [] | |
| current_image = None # To store the uploaded image | |
| def process_image(question): | |
| global current_image | |
| if current_image is None: | |
| return "Please upload an image first." | |
| image_np = np.load(current_image) # Load the stored .npy image | |
| image_tokens = "<im_patch>" * proj_out_num | |
| input_txt = image_tokens + question | |
| input_id = tokenizer(input_txt, return_tensors="pt")['input_ids'].to(device=device) | |
| # Prepare image for model | |
| image_pt = torch.from_numpy(image_np).unsqueeze(0).to(dtype=dtype, device=device) | |
| # Generate response | |
| generation = model.generate(image_pt, input_id, max_new_tokens=256, do_sample=True, top_p=0.9, temperature=1.0) | |
| generated_texts = tokenizer.batch_decode(generation, skip_special_tokens=True) | |
| return generated_texts[0] | |
| # Function to update chat | |
| def chat_interface(question): | |
| global chat_history | |
| response = process_image(question) | |
| chat_history.append((question, response)) | |
| return chat_history | |
| # Function to handle image upload | |
| def upload_image(image): | |
| global current_image | |
| current_image = image.name | |
| return "Image uploaded successfully!" | |
| # Gradio UI | |
| with gr.Blocks(theme=gr.themes.Soft()) as chat_ui: | |
| gr.Markdown("# 🏥 Medical Image Analysis Chatbot") | |
| with gr.Row(): | |
| with gr.Column(scale=1, min_width=200): | |
| chat_list = gr.Chatbot(value=[], label="Chat History", elem_id="chat-history") | |
| with gr.Column(scale=4): | |
| uploaded_image = gr.File(label="Upload .npy Image", type="filepath") | |
| upload_status = gr.Textbox(label="Status", interactive=False) | |
| question_input = gr.Textbox(label="Ask a question", placeholder="Ask something about the image...") | |
| submit_button = gr.Button("Send") | |
| uploaded_image.upload(upload_image, uploaded_image, upload_status) | |
| submit_button.click(chat_interface, question_input, chat_list) | |
| question_input.submit(chat_interface, question_input, chat_list) | |
| chat_ui.launch() |