Update app.py

app.py

@@ -5,129 +5,112 @@ from PIL import Image
 from datasets import load_dataset
 from transformers import TextStreamer
 import matplotlib.pyplot as plt
-
 import gradio as gr
+import random
+import numpy as np
 
-# Load the model
-model, tokenizer = FastVisionModel.from_pretrained(
-    "0llheaven/Llama-3.2-11B-Vision-Radiology-mini",
-    load_in_4bit=True,
-    use_gradient_checkpointing="unsloth",
-).to("cpu")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# เปลี่ยนโหมดของโมเดลเป็นสำหรับ inference
-FastVisionModel.for_inference(model)
+def set_seed(seed_value=42):
+    random.seed(seed_value)
+    np.random.seed(seed_value)
+    torch.manual_seed(seed_value)
+    torch.cuda.manual_seed_all(seed_value)
+    
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
 
-# ตัวแปรสำหรับแคช
-cached_image = None
-cached_response = None
+set_seed(42)
 
-# ฟังก์ชันประมวลผลภาพและสร้างคำอธิบาย
-def predict_radiology_description(image, instruction):
-    global cached_image, cached_response
-
-    try:
+model, tokenizer = FastVisionModel.from_pretrained(
+    "0llheaven/llama-3.2-11B-Vision-Instruct-Finetune",
+    load_in_4bit = True,
+    use_gradient_checkpointing = "unsloth",
+)
 
-        current_image_tensor = torch.tensor(image.getdata())
+FastVisionModel.for_inference(model)
 
-        # ตรวจสอบว่าภาพเหมือนเดิมและข้อความเหมือนเดิมหรือไม่
-        if cached_image is not None and torch.equal(cached_image, current_image_tensor):
-            # ใช้ cached_response กับ text ใหม่
-            return cached_response
+instruction = "You are an expert radiographer. Describe accurately what you see in this image."
 
-        # เตรียมข้อความในรูปแบบที่โมเดลรองรับ
+def predict_radiology_description(image, temperature, use_top_p, top_p_value, use_min_p, min_p_value):
+    try:
+        set_seed(42)
         messages = [{"role": "user", "content": [
             {"type": "image"},
             {"type": "text", "text": instruction}
         ]}]
         input_text = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
 
-        # เตรียม input สำหรับโมเดล
         inputs = tokenizer(
             image,
             input_text,
             add_special_tokens=False,
             return_tensors="pt",
-        ).to("cpu")
+        ).to("cuda")
 
-        # ใช้ TextStreamer สำหรับการพยากรณ์
         text_streamer = TextStreamer(tokenizer, skip_prompt=True)
 
-        # ทำนายข้อความ
+        generate_kwargs = {
+            "max_new_tokens": 512,
+            "use_cache": True,
+            "temperature": temperature,
+        }
+        if use_top_p:
+            generate_kwargs["top_p"] = top_p_value
+        if use_min_p:
+            generate_kwargs["min_p"] = min_p_value
+
         output_ids = model.generate(
             **inputs,
             streamer=text_streamer,
-            max_new_tokens=256,
-            use_cache=True,
-            temperature=1.5,
-            min_p=0.1
+            **generate_kwargs
         )
 
-        # แปลงข้อความที่สร้างเป็นผลลัพธ์
         generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+        return generated_text.replace("assistant", "\n\nassistant").strip()
 
-        cached_image = current_image_tensor  # แคชภาพเป็น Tensor
-        cached_response = generated_text.replace("assistant", "\n\nAssistant").strip()
-        return cached_response
     except Exception as e:
         return f"Error: {str(e)}"
 
-# ฟังก์ชัน ChatBot
-def chat_process(image, instruction, history=None):
-    if history is None:
-        history = []
-
-    # ประมวลผลภาพและคำสั่ง
-    response = predict_radiology_description(image, instruction)
+with gr.Blocks() as interface:
+    gr.Markdown("<h1><center>Radiology Image Description Generator</center></h1>")
+    gr.Markdown("Upload a radiology image, adjust temperature and top-p, and the model will describe the findings in the image")
+    with gr.Row(): 
+        with gr.Column():
+            image_input = gr.Image(type="pil", label="Upload")
+        with gr.Column():
+            output_text = gr.Textbox(label="Generated Description")
 
-    # อัปเดตประวัติ
-    history.append((instruction, response))
-    return history, history
+    with gr.Row():
+        with gr.Column(scale=0.5):
+            temperature_slider = gr.Slider(0.1, 2.0, step=0.1, value=1.0, label="temperature")
 
-import warnings
-warnings.filterwarnings("ignore", category=UserWarning, module="gradio.helpers")
+            use_top_p_checkbox = gr.Checkbox(label="Use top-p", value=True)
+            top_p_slider = gr.Slider(0.1, 1.0, step=0.05, value=0.9, label="top-p")
 
-# UI ของ Gradio
-with gr.Blocks() as demo:
-    gr.Markdown("# 🩻 Radiology Image ChatBot")
-    gr.Markdown("Upload a radiology image and provide an instruction for the AI to describe the findings.")
-    gr.Markdown("Example instruction : You are an expert radiographer. Describe accurately what you see in this image.")
+            use_min_p_checkbox = gr.Checkbox(label="Use min-p", value=False)
+            min_p_slider = gr.Slider(0.0, 1.0, step=0.05, value=0.1, label="min-p", visible=False)
 
-    with gr.Row():
-        with gr.Column():
-            # อัปโหลดรูปภาพ
-            image_input = gr.Image(type="pil", label="Upload Radiology Image")
-            # ป้อนคำสั่ง (instruction)
-            instruction_input = gr.Textbox(
-                label="Instruction",
-                value="You are an expert radiographer. Describe accurately what you see in this image.",
-                placeholder="Provide specific instructions..."
+            # Update visibility of sliders
+            use_top_p_checkbox.change(
+                lambda use_top_p: gr.update(visible=use_top_p),
+                inputs=use_top_p_checkbox,
+                outputs=top_p_slider
+            )
+            use_min_p_checkbox.change(
+                lambda use_min_p: gr.update(visible=use_min_p),
+                inputs=use_min_p_checkbox,
+                outputs=min_p_slider
             )
-        with gr.Column():
-            # แสดงประวัติ Chat
-            chatbot = gr.Chatbot(label="Chat History")
 
-    with gr.Row():
-        clear_btn = gr.Button("Clear")
-        submit_btn = gr.Button("Submit")
-
-    # การทำงานของปุ่ม Submit พร้อมล้างเฉพาะข้อความใน instruction_input
-    submit_btn.click(
-        lambda image, instruction, history: (
-            *chat_process(image, instruction, history),
-            image,  # รีเซ็ตค่า image_input
-            ""
-        ),
-        inputs=[image_input, instruction_input, chatbot],
-        outputs=[chatbot, chatbot, image_input, instruction_input]
-    )
+    generate_button = gr.Button("Generate Description")
 
-    # การทำงานของปุ่ม Clear
-    clear_btn.click(
-        lambda: (None, None, None, None),
-        inputs=[],
-        outputs=[chatbot, chatbot, image_input, instruction_input]
+    # Link function to UI
+    generate_button.click(
+        predict_radiology_description,
+        inputs=[image_input, temperature_slider, use_top_p_checkbox, top_p_slider, use_min_p_checkbox, min_p_slider],
+        outputs=output_text
     )
 
-# รันแอป
-demo.launch(debug=True)
+# Gradio
+interface.launch(share=True, debug=True)