Create DocBrain.py

DocBrain.py

@@ -0,0 +1,240 @@
+import gradio as gr
+from transformers import AutoProcessor, AutoModelForVision2Seq
+import torch
+from PaitentVoiceToText import record_and_transcribe  # Your STT function
+from DocVoice import text_to_speech  # Your TTS function
+
+# -------------------
+# 1️⃣ Load Model & Processor
+# -------------------
+def load_model():
+    local_dir = r"C:\Users\JAY\Downloads\model\CHATDOCMODEL"
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    dtype = torch.float16 if device == "cuda" else torch.float32
+
+    processor = AutoProcessor.from_pretrained(local_dir, trust_remote_code=True)
+    model = AutoModelForVision2Seq.from_pretrained(
+        local_dir,
+        dtype=dtype,
+        device_map=None
+    )
+    model.to(device)
+    return processor, model, device
+
+processor, model, device = load_model()
+
+# -------------------
+# 2️⃣ Chat Logic Functions
+# -------------------
+def process_message(message, history, question_count):
+    if not message.strip():
+        return history, history, question_count
+    
+    history.append([message, None])
+    question_count += 1
+    
+    should_analyze = (
+        question_count >= 6 or 
+        any(word in message.lower() for word in ["analysis", "diagnose", "what do you think", "causes"])
+    )
+
+    if should_analyze:
+        system_prompt = (
+            "You are a medical doctor. Based on the patient's responses, provide a comprehensive analysis "
+            "of potential causes for their symptoms. Start with 'Based on the information provided by the patient, "
+            "potential causes of [symptoms] could include:' and list 3-4 possible diagnoses with brief explanations. "
+            "Format as numbered list with diagnosis name and short explanation."
+        )
+    else:
+        system_prompt = (
+            "You are a medical doctor conducting a patient interview. Ask ONE specific, direct medical question "
+            "to gather important diagnostic information. Keep it brief - just ask the question without explanations. "
+            "Focus on key areas like: age, medical history, medications, lifestyle, family history, or symptom details."
+        )
+    
+    dialogue = []
+    for user_msg, bot_msg in history[:-1]:
+        if user_msg:
+            dialogue.append(f"Patient: {user_msg}")
+        if bot_msg:
+            dialogue.append(f"Doctor: {bot_msg}")
+    dialogue.append(f"Patient: {message}")
+    
+    conversation = "\n".join(dialogue)
+    prompt = f"{system_prompt}\n\nConversation:\n{conversation}\nDoctor:"
+
+    inputs = processor(text=prompt, images=None, return_tensors="pt").to(device)
+    max_tokens = 1000 if should_analyze else 25
+    
+    with torch.inference_mode():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=max_tokens,
+            do_sample=True,
+            temperature=0.6,
+            top_p=0.9,
+            repetition_penalty=1.1,
+            pad_token_id=processor.tokenizer.eos_token_id,
+        )
+
+    input_length = inputs["input_ids"].shape[1]
+    generated_tokens = outputs[:, input_length:]
+    response = processor.batch_decode(generated_tokens, skip_special_tokens=True)[0].strip()
+
+    if response.lower().startswith("doctor:"):
+        response = response[7:].strip()
+
+    if not should_analyze:
+        sentences = response.split('?')
+        if len(sentences) > 1:
+            response = sentences[0].strip() + '?'
+        cleanup_starts = [
+            "I need to ask",
+            "Let me ask",
+            "I would like to know",
+            "Can you tell me",
+            "It would help if",
+        ]
+        for phrase in cleanup_starts:
+            if response.startswith(phrase):
+                parts = response.split(',', 1)
+                if len(parts) > 1:
+                    response = parts[1].strip()
+                    if not response.endswith('?'):
+                        response += '?'
+
+    history[-1][1] = response
+    if should_analyze:
+        question_count = 0
+    
+    return history, history, question_count
+
+def force_analysis(history, question_count):
+    return history, 10
+
+def clear_chat():
+    return [], [], 0
+
+# -------------------
+# 3️⃣ TTS Helper
+# -------------------
+def play_assistant_audio(response_text):
+    if response_text:
+        text_to_speech(response_text)
+    return None
+
+# -------------------
+# 4️⃣ Gradio Interface
+# -------------------
+with gr.Blocks(title="ChatDOC", theme=gr.themes.Soft()) as demo:
+    question_count_state = gr.State(0)
+    assistant_responses_state = gr.State([])
+
+    gr.Markdown(
+        """
+        # 🩺 Chat with ChatDOC
+        Welcome! I'm your AI medical assistant. Please describe your symptoms and I'll ask relevant questions to help understand your condition better.
+        """
+    )
+    
+    chatbot = gr.Chatbot(
+        value=[],
+        height=400,
+        show_label=False,
+        avatar_images=(
+            r"C:\Users\JAY\Downloads\model\user_msg.png",
+            r"C:\Users\JAY\Downloads\model\bot_msg.jpg"
+        ),
+        bubble_full_width=False
+    )
+    
+    with gr.Row():
+        msg = gr.Textbox(
+            placeholder="Describe your symptoms...",
+            scale=4,
+            container=False,
+            show_label=False
+        )
+        send_btn = gr.Button("Send", variant="primary", scale=1)
+        mic_btn = gr.Button("🎤 Speak", variant="secondary", scale=1)
+    
+    with gr.Row():
+        analysis_btn = gr.Button("Request Analysis", variant="secondary")
+        clear_btn = gr.Button("Clear Chat", variant="stop")
+        play_audio_btn = gr.Button("🔊 Play Assistant Response", variant="secondary")
+    
+    # -------------------
+    # Update assistant responses
+    # -------------------
+    def update_assistant_responses(history, assistant_responses):
+        if history and history[-1][1]:
+            assistant_responses.append(history[-1][1])
+        return assistant_responses
+
+    # -------------------
+    # Submit handlers
+    # -------------------
+    def user_submit(message, history, question_count, assistant_responses):
+        history, updated_history, question_count = process_message(message, history, question_count)
+        assistant_responses = update_assistant_responses(history, assistant_responses)
+        return updated_history, updated_history, question_count, assistant_responses
+
+    def mic_submit(history, question_count, assistant_responses):
+        user_text = record_and_transcribe(duration=5)
+        # Show user message immediately
+        history.append([user_text, None])
+        history, updated_history, question_count = process_message(user_text, history, question_count)
+        assistant_responses = update_assistant_responses(history, assistant_responses)
+        return updated_history, updated_history, question_count, assistant_responses
+
+    def clear_input():
+        return ""
+    
+    # -------------------
+    # Connect buttons
+    # -------------------
+    send_btn.click(
+        user_submit,
+        inputs=[msg, chatbot, question_count_state, assistant_responses_state],
+        outputs=[chatbot, chatbot, question_count_state, assistant_responses_state]
+    ).then(clear_input, outputs=[msg])
+    
+    msg.submit(
+        user_submit,
+        inputs=[msg, chatbot, question_count_state, assistant_responses_state],
+        outputs=[chatbot, chatbot, question_count_state, assistant_responses_state]
+    ).then(clear_input, outputs=[msg])
+    
+    mic_btn.click(
+        mic_submit,
+        inputs=[chatbot, question_count_state, assistant_responses_state],
+        outputs=[chatbot, chatbot, question_count_state, assistant_responses_state]
+    )
+    
+    analysis_btn.click(
+        force_analysis,
+        inputs=[chatbot, question_count_state],
+        outputs=[chatbot, question_count_state]
+    )
+    
+    clear_btn.click(
+        clear_chat,
+        outputs=[chatbot, chatbot, question_count_state]
+    )
+
+    play_audio_btn.click(
+        lambda assistant_responses: play_assistant_audio(assistant_responses[-1]) if assistant_responses else None,
+        inputs=[assistant_responses_state],
+        outputs=[]
+    )
+
+# -------------------
+# 5️⃣ Launch
+# -------------------
+if __name__ == "__main__":
+    demo.launch(
+        server_name="127.0.0.1",
+        server_port=7860,
+        share=False,
+        debug=True
+    )