Debug

app.py

@@ -98,41 +98,26 @@ def chat_with_model(messages):
         return
 
     current_id = patient_id.value
-    if current_id is None:
+    if not current_id:
         yield messages
         return
 
-    # 🛠 Missing variable initializations
     max_new_tokens = 1024
     output_text = ""
     in_think = False
     generated_tokens = 0
 
-    pad_id = current_tokenizer.pad_token_id
+    pad_id = current_tokenizer.pad_token_id or current_tokenizer.unk_token_id or 0
     eos_id = current_tokenizer.eos_token_id
-    if pad_id is None:
-        pad_id = current_tokenizer.unk_token_id or 0
-
-    # Remove the initial welcome if present
-    filtered_messages = [msg for msg in messages if not (msg["role"] == "assistant" and "Welcome to the Radiologist's Companion" in msg["content"])]
 
-    # Build system context
+    # --- Build system context
     system_messages = [
         {
             "role": "system",
             "content": (
                 "You are a radiologist's companion, here to answer questions about the patient and assist in the diagnosis if asked to do so. "
                 "You are able to call specialized tools. "
-                "At the moment, you have one tool available: an organ segmentation algorithm for abdominal CTs.\n\n"
-                "If the user requests an organ segmentation, output a JSON object in this structure:\n"
-                "{\n"
-                "  \"function\": \"segment_organ\",\n"
-                "  \"arguments\": {\n"
-                "    \"scan_path\": \"<path_to_ct_scan>\",\n"
-                "    \"organ\": \"<organ_name>\"\n"
-                "  }\n"
-                "}\n\n"
-                "Once you call the function, the app will execute it and return the result."
+                "At the moment, you have one tool available: an organ segmentation algorithm for abdominal CTs."
             )
         },
         {
@@ -141,8 +126,13 @@ def chat_with_model(messages):
         }
     ]
 
+    # Remove welcome message (only once shown)
+    # filtered_messages = [msg for msg in messages if not (msg["role"] == "assistant" and "Welcome to the Radiologist's Companion" in msg["content"])]
+
+    # FULL conversation
     full_messages = system_messages + filtered_messages
 
+    # --- Generate from full context
     prompt = format_prompt(full_messages)
 
     device = torch.device("cuda")
@@ -169,15 +159,13 @@ def chat_with_model(messages):
     thread = threading.Thread(target=current_model.generate, kwargs=generation_kwargs)
     thread.start()
 
-    messages = full_messages.copy()
-    messages.append({"role": "assistant", "content": ""})
-
-    print(messages)
+    # Now extend previous messages
+    updated_messages = messages.copy()
+    updated_messages.append({"role": "assistant", "content": ""})
 
     for token_info in streamer:
         token_str = token_info["token"]
         token_id = token_info["token_id"]
-        is_special = token_info["is_special"]
 
         if token_id == eos_id:
             break
@@ -196,27 +184,26 @@ def chat_with_model(messages):
 
         if "\nUser" in output_text:
             output_text = output_text.split("\nUser")[0].rstrip()
-            messages[-1]["content"] = output_text
+            updated_messages[-1]["content"] = output_text
             break
 
         generated_tokens += 1
         if generated_tokens >= max_new_tokens:
             break
 
-        messages[-1]["content"] = output_text
+        updated_messages[-1]["content"] = output_text
 
-        # Save conversation per patient
-        patient_conversations[current_id] = messages
-
-        yield messages
+        patient_conversations[current_id] = updated_messages  
+        yield updated_messages
 
     if in_think:
         output_text += "*"
-        messages[-1]["content"] = output_text
+        updated_messages[-1]["content"] = output_text
 
     torch.cuda.empty_cache()
-    messages[-1]["content"] = output_text
-    return messages
+    updated_messages[-1]["content"] = output_text
+    return updated_messages
+
 
 
 def load_model_on_selection(model_name, progress=gr.Progress(track_tqdm=False)):