medical_chatbot.py

# Setup and Installation

import torch
print("🖥️ System Check:")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU device: {torch.cuda.get_device_name(0)}")
    print(f"GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
else:
    print("⚠️ No GPU detected - BioGPT will run on CPU (much slower)")

print("\n🔧 Installing required packages...")

# Import Libraries

import os
import re
import torch
import warnings
import numpy as np
import faiss  # FAISS for vector search
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    pipeline,
    BitsAndBytesConfig
)
from sentence_transformers import SentenceTransformer
from typing import List, Dict, Optional
import time
from datetime import datetime
import json
import pickle

# Suppress warnings for cleaner output
warnings.filterwarnings('ignore')

print("📚 Libraries imported successfully!")
print(f"🔍 FAISS version: {faiss.__version__}")
print("🎯 Using FAISS for vector search (ChromaDB completely removed)")

# File Upload Helper

from google.colab import files
import io

def upload_medical_data():
    """Upload your Pediatric_cleaned.txt file"""
    print("📁 Please upload your Pediatric_cleaned.txt file:")
    uploaded = files.upload()

    # Get the uploaded file
    filename = list(uploaded.keys())[0]
    print(f"✅ File '{filename}' uploaded successfully!")

    # Read the content
    content = uploaded[filename].decode('utf-8')

    # Save it locally in Colab
    with open('Pediatric_cleaned.txt', 'w', encoding='utf-8') as f:
        f.write(content)

    print(f"📝 File saved as 'Pediatric_cleaned.txt' ({len(content)} characters)")
    return 'Pediatric_cleaned.txt'

medical_file = 'Pediatric_cleaned.txt'

# BioGPT Medical Chatbot Class

from typing import List, Dict, Optional # Import List, Dict, Optional

class ColabBioGPTChatbot:
    def setup_biogpt(self):
        """Setup BioGPT model with fallback to base BioGPT if Large fails"""
        print("🧠 Attempting to load BioGPT-Large...")
        model_name = "microsoft/BioGPT-Large"

        try:
            if self.use_8bit:
                quantization_config = BitsAndBytesConfig(
                    load_in_8bit=True,
                    llm_int8_threshold=6.0,
                    llm_int8_has_fp16_weight=False,
                )
            else:
                quantization_config = None

            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
            if self.tokenizer.pad_token is None:
                self.tokenizer.pad_token = self.tokenizer.eos_token

            self.model = AutoModelForCausalLM.from_pretrained(
                model_name,
                quantization_config=quantization_config,
                torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
                device_map="auto" if self.device == "cuda" else None,
                trust_remote_code=True
            )

            if self.device == "cuda" and quantization_config is None:
                self.model = self.model.to(self.device)

            print("✅ BioGPT-Large loaded successfully!")
            self.test_biogpt()

        except Exception as e:
            print(f"❌ BioGPT-Large loading failed: {e}")
            print("🔁 Falling back to base model: microsoft/BioGPT")
            self.setup_fallback_biogpt()

        def fallback_biogpt():
    """Fallback to microsoft/BioGPT if BioGPT-Large fails"""
    print("Loading fallback model: microsoft/BioGPT")
    # You can implement the actual fallback loading here later
    pass

        model_name = "microsoft/BioGPT"

        try:
            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
            if self.tokenizer.pad_token is None:
                self.tokenizer.pad_token = self.tokenizer.eos_token

            self.model = AutoModelForCausalLM.from_pretrained(
                model_name,
                torch_dtype=torch.float32,
                trust_remote_code=True
            )
            self.model = self.model.to(self.device)

            print("✅ Base BioGPT model loaded as fallback.")
            self.test_biogpt()

        except Exception as e:
            print(f"❌ Failed to load fallback BioGPT: {e}")
            self.model = None
            self.tokenizer = None

(self):
        if self.tokenizer is None or self.model is None:
            print("⚠️ Skipping test — model or tokenizer not available.")
            return
        print("🧪 Testing BioGPT...")
        try:
            input_text = "What is pneumonia?"
            inputs = self.tokenizer(input_text, return_tensors="pt").to(self.device)
            outputs = self.model.generate(**inputs, max_new_tokens=30)
            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
            print(f"🧠 Model output: {response}")
        except Exception as e:
            print(f"⚠️ BioGPT test failed: {e}")

        def __init__(self, use_gpu=True, use_8bit=True):
            """Initialize BioGPT chatbot optimized for Google Colab"""
        print("🏥 Initializing...")
        self.use_gpu = use_gpu
        self.use_8bit = use_8bit
        self.device = "cuda" if torch.cuda.is_available() and use_gpu else "cpu"
        print(f"🖥️ System Check:\nCUDA available: {torch.cuda.is_available()}\nUsing device: {self.device}")
        self.tokenizer = None
        self.model = None
        self.vectorstore = None
        self.history = []
        self.embeddings = None

        print("🔧 Installing required packages...")

        try:
                import transformers
                import faiss
                print("📚 Libraries imported successfully!")
                print(f"🔍 FAISS version: {faiss.__version__}")
        except ImportError as e:
                print(f"❌ Missing package: {e}")

                self.setup_biogpt()
                self.load_sentence_transformer()
                self.vectorstore = None

        def setup_biogpt(self):
            """Setup BioGPT model with fallback to base BioGPT if Large fails"""
            print("🧠 Attempting to load BioGPT-Large...")
            model_name = "microsoft/BioGPT-Large"
        try:
            if self.use_8bit and self.device == "cuda":
                from transformers import BitsAndBytesConfig
                bnb_config = BitsAndBytesConfig(load_in_8bit=True)
                self.model = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=bnb_config, device_map="auto")
            else:
                self.model = AutoModelForCausalLM.from_pretrained(model_name).to(self.device)

            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
            print("✅ BioGPT-Large loaded successfully.")
        except Exception as e:
            print(f"❌ BioGPT-Large loading failed: {e}")
            print("🔁 Falling back to base model: microsoft/BioGPT")
            self.setup_fallback_biogpt()

        def setup_fallback_biogpt(self):
            """Fallback to microsoft/BioGPT if BioGPT-Large fails"""
            print("Loading fallback model: microsoft/BioGPT")
            model_name = "microsoft/BioGPT"
            self.model = AutoModelForCausalLM.from_pretrained(model_name).to(self.device)
            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
            print("✅ Fallback BioGPT model loaded.")

        def test_biogpt(self):
            """Test BioGPT with a simple medical query"""
            print("🧪 Testing BioGPT...")
            if not self.model or not self.tokenizer:
                print("⚠️ BioGPT not properly initialized.")
                return

            input_text = "What are symptoms of measles?"
            inputs = self.tokenizer(input_text, return_tensors="pt").to(self.device)
            outputs = self.model.generate(**inputs, max_new_tokens=50)
            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
            print(f"💬 Test response: {response}")


        def load_medical_data(self, file_path):
            """Load and process medical data with progress tracking"""
            print(f"📖 Loading medical data from {file_path}...")

            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    text = f.read()
                print(f"📄 File loaded: {len(text):,} characters")
            except FileNotFoundError:
                print(f"❌ File {file_path} not found!")
                raise ValueError("Failed to load medical data")

            # Create chunks optimized for medical content
            print("📝 Creating medical-optimized chunks...")
            chunks = self.create_medical_chunks(text)
            print(f"📋 Created {len(chunks)} medical chunks")

            self.knowledge_chunks = chunks

            # Generate embeddings with progress and add to FAISS index
            if self.use_embeddings and self.embedding_model and self.faiss_ready:
                return self.generate_embeddings_with_progress(chunks)

            print("✅ Medical data loaded (text search mode)")
            return

    def create_medical_chunks(self, text: str, chunk_size: int = 400) -> List[Dict]:
        """Create medically-optimized text chunks"""
        chunks = []

        # Split by medical sections first
        medical_sections = self.split_by_medical_sections(text)

        chunk_id = 0
        for section in medical_sections:
            if len(section.split()) > chunk_size:
                # Split large sections by sentences
                sentences = re.split(r'[.!?]+', section)
                current_chunk = ""

                for sentence in sentences:
                    sentence = sentence.strip()
                    if not sentence:
                        continue

                    if len(current_chunk.split()) + len(sentence.split()) < chunk_size:
                        current_chunk += sentence + ". "
                    else:
                        if current_chunk.strip():
                            chunks.append({
                                'id': chunk_id,
                                'text': current_chunk.strip(),
                                'medical_focus': self.identify_medical_focus(current_chunk)
                            })
                            chunk_id += 1
                        current_chunk = sentence + ". "

                if current_chunk.strip():
                    chunks.append({
                        'id': chunk_id,
                        'text': current_chunk.strip(),
                        'medical_focus': self.identify_medical_focus(current_chunk)
                    })
                    chunk_id += 1
            else:
                chunks.append({
                    'id': chunk_id,
                    'text': section,
                    'medical_focus': self.identify_medical_focus(section)
                })
                chunk_id += 1

        return chunks

    def split_by_medical_sections(self, text: str) -> List[str]:
        """Split text by medical sections"""
        # Look for medical section headers
        section_patterns = [
            r'\n\s*(?:SYMPTOMS?|TREATMENT|DIAGNOSIS|CAUSES?|PREVENTION|MANAGEMENT).*?\n',
            r'\n\s*\d+\.\s+',  # Numbered sections
            r'\n\n+'  # Paragraph breaks
        ]

        sections = [text]
        for pattern in section_patterns:
            new_sections = []
            for section in sections:
                splits = re.split(pattern, section, flags=re.IGNORECASE)
                new_sections.extend([s.strip() for s in splits if len(s.strip()) > 100])
            sections = new_sections

        return sections

    def identify_medical_focus(self, text: str) -> str:
        """Identify the medical focus of a text chunk"""
        text_lower = text.lower()

        # Medical categories
        categories = {
            'pediatric_symptoms': ['fever', 'cough', 'rash', 'vomiting', 'diarrhea'],
            'treatments': ['treatment', 'therapy', 'medication', 'antibiotics'],
            'diagnosis': ['diagnosis', 'diagnostic', 'symptoms', 'signs'],
            'emergency': ['emergency', 'urgent', 'serious', 'hospital'],
            'prevention': ['prevention', 'vaccine', 'immunization', 'avoid']
        }

        for category, keywords in categories.items():
            if any(keyword in text_lower for keyword in keywords):
                return category

        return 'general_medical'

    def generate_embeddings_with_progress(self, chunks: List[Dict]) -> bool:
        """Generate embeddings with progress tracking and add to FAISS index"""
        print("🔮 Generating medical embeddings and adding to FAISS index...")

        if not self.embedding_model or not self.faiss_index:
            print("❌ Embedding model or FAISS index not available.")
            raise ValueError("Failed to load medical data")

        try:
            texts = [chunk['text'] for chunk in chunks]

            # Generate embeddings in batches with progress
            batch_size = 32
            all_embeddings = []

            for i in range(0, len(texts), batch_size):
                batch_texts = texts[i:i+batch_size]
                batch_embeddings = self.embedding_model.encode(batch_texts, show_progress_bar=False)
                all_embeddings.extend(batch_embeddings)

                # Show progress
                progress = min(i + batch_size, len(texts))
                print(f"   Progress: {progress}/{len(texts)} chunks processed", end='\r')

            print(f"\n   ✅ Generated embeddings for {len(texts)} chunks")

            # Add embeddings to FAISS index
            print("💾 Adding embeddings to FAISS index...")
            self.faiss_index.add(np.array(all_embeddings))

            print("✅ Medical embeddings added to FAISS index successfully!")
            return True

        except Exception as e:
            print(f"❌ Embedding generation or FAISS add failed: {e}")
            raise ValueError("Failed to load medical data")


    def retrieve_medical_context(self, query: str, n_results: int = 3) -> List[str]:
        """Retrieve relevant medical context using embeddings or keyword search"""
        if self.use_embeddings and self.embedding_model and self.faiss_ready:
            try:
                # Generate query embedding
                query_embedding = self.embedding_model.encode([query])

                # Search for similar content in FAISS index
                distances, indices = self.faiss_index.search(np.array(query_embedding), n_results)

                # Retrieve the corresponding chunks
                context_chunks = [self.knowledge_chunks[i]['text'] for i in indices[0] if i != -1]

                if context_chunks:
                    return context_chunks

            except Exception as e:
                print(f"⚠️ Embedding search failed: {e}")

        # Fallback to keyword search
        print("⚠️ Falling back to keyword search.")
        return self.keyword_search_medical(query, n_results)


    def keyword_search_medical(self, query: str, n_results: int) -> List[str]:
        """Medical-focused keyword search"""
        if not self.knowledge_chunks:
            return []

        query_words = set(query.lower().split())
        chunk_scores = []

        for chunk_info in self.knowledge_chunks:
            chunk_text = chunk_info['text']
            chunk_words = set(chunk_text.lower().split())

            # Calculate relevance score
            word_overlap = len(query_words.intersection(chunk_words))
            base_score = word_overlap / len(query_words) if query_words else 0

            # Boost medical content
            medical_boost = 0
            if chunk_info.get('medical_focus') in ['pediatric_symptoms', 'treatments', 'diagnosis']:
                medical_boost = 0.5

            final_score = base_score + medical_boost

            if final_score > 0:
                chunk_scores.append((final_score, chunk_text))

        # Return top matches
        chunk_scores.sort(reverse=True)
        return [chunk for _, chunk in chunk_scores[:n_results]]

    def generate_biogpt_response(self, context: str, query: str) -> str:
        """Generate medical response using BioGPT"""
        if not self.model or not self.tokenizer:
            return "Medical model not available. Please check the setup."

        try:
            # Create medical-focused prompt
            prompt = f"""Medical Context: {context[:800]}

Question: {query}

Medical Answer:"""

            # Tokenize input
            inputs = self.tokenizer(
                prompt,
                return_tensors="pt",
                truncation=True,
                max_length=1024
            )

            # Move inputs to the correct device
            if self.device == "cuda":
                inputs = {k: v.to(self.device) for k, v in inputs.items()}

            # Generate response
            with torch.no_grad():
                outputs = self.model.generate(
                    **inputs,
                    max_new_tokens=150,
                    do_sample=True,
                    temperature=0.7,
                    top_p=0.9,
                    pad_token_id=self.tokenizer.eos_token_id,
                    repetition_penalty=1.1
                )

            # Decode response
            full_response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)

            # Extract just the generated part
            if "Medical Answer:" in full_response:
                generated_response = full_response.split("Medical Answer:")[-1].strip()
            else:
                generated_response = full_response[len(prompt):].strip()

            # Clean up response
            cleaned_response = self.clean_medical_response(generated_response)

            return cleaned_response

        except Exception as e:
            print(f"⚠️ BioGPT generation failed: {e}")
            return self.fallback_response(context, query)

    def clean_medical_response(self, response: str) -> str:
        """Clean and format medical response"""
        # Remove incomplete sentences and limit length
        sentences = re.split(r'[.!?]+', response)
        clean_sentences = []

        for sentence in sentences:
            sentence = sentence.strip()
            if len(sentence) > 10 and not sentence.endswith(('and', 'or', 'but', 'however')):
                clean_sentences.append(sentence)
            if len(clean_sentences) >= 3:  # Limit to 3 sentences
                break

        if clean_sentences:
            cleaned = '. '.join(clean_sentences) + '.'
        else:
            cleaned = response[:200] + '...' if len(response) > 200 else response

        return cleaned

    def fallback_response(self, context: str, query: str) -> str:
        """Fallback response when BioGPT fails"""
        # Extract key sentences from context
        sentences = [s.strip() for s in context.split('.') if len(s.strip()) > 20]

        if sentences:
            response = sentences[0] + '.'
            if len(sentences) > 1:
                response += ' ' + sentences[1] + '.'
        else:
            response = context[:300] + '...'

        return response

    def handle_conversational_interactions(self, query: str) -> Optional[str]:
        """Handle comprehensive conversational interactions"""
        query_lower = query.lower().strip()

        # Use more specific patterns for greetings
        greeting_patterns = [
            r'^\s*(hello|hi|hey|hiya|howdy)\s*$',
            r'^\s*(good morning|good afternoon|good evening|good day)\s*$',
            r'^\s*(what\'s up|whats up|sup|yo)\s*$',
            r'^\s*(greetings|salutations)\s*$',
            r'^\s*(how are you|how are you doing|how\'s it going|hows it going)\s*$',
            r'^\s*(good to meet you|nice to meet you|pleased to meet you)\s*$'
        ]

        for pattern in greeting_patterns:
            if re.match(pattern, query_lower):
                responses = [
                    "👋 Hello! I'm BioGPT, your professional medical AI assistant specialized in pediatric medicine. I'm here to provide evidence-based medical information. What health concern can I help you with today?",
                    "🏥 Hi there! I'm a medical AI assistant powered by BioGPT, trained on medical literature. I can help answer questions about children's health and medical conditions. How can I assist you?",
                    "👋 Greetings! I'm your AI medical consultant, ready to help with pediatric health questions using the latest medical knowledge. What would you like to know about?"
                ]
                return np.random.choice(responses)

        # ===== THANKS & APPRECIATION =====
        thanks_patterns = [
            ['thank you', 'thanks', 'thx', 'ty', 'thank you so much', 'thanks a lot', 'much appreciated', 'really appreciate it', 'i appreciate it', 'grateful', 'that was helpful', 'very helpful', 'awesome', 'perfect', 'great', 'excellent', 'wonderful', 'that helped', 'exactly what i needed', 'very informative', 'good information']
        ]

        for pattern_group in thanks_patterns:
            if any(keyword in query_lower for keyword in pattern_group):
                responses = [
                    "🙏 You're very welcome! I'm glad I could provide helpful medical information. Remember, this is educational guidance - always consult your healthcare provider for personalized medical advice. Feel free to ask more questions!",
                    "😊 Happy to help! Providing accurate medical information is what I'm here for. If you have any other pediatric health questions, don't hesitate to ask.",
                    "🤗 You're most welcome! I'm pleased the medical information was useful. Please remember to consult with healthcare professionals for any medical decisions. What else can I help you with?"
                ]
                return np.random.choice(responses)

        # ===== GOODBYES =====
        goodbye_patterns = [
            ['bye', 'goodbye', 'farewell', 'see you', 'later', 'see ya', 'catch you later', 'talk to you later', 'ttyl', 'have a good day', 'have a great day', 'take care', 'until next time', 'i need to go', 'that\'s all for now', 'no more questions']
        ]

        for pattern_group in goodbye_patterns:
            if any(keyword in query_lower for keyword in pattern_group):
                responses = [
                    "👋 Goodbye! Take excellent care of yourself and your little ones. Remember, I'm here whenever you need reliable pediatric medical information. Stay healthy! 🏥",
                    "🌟 Farewell! Wishing you and your family good health. Don't hesitate to return if you have more medical questions. Take care!",
                    "👋 See you later! Hope the medical information was helpful. Remember to always consult healthcare professionals for medical decisions. Stay well!"
                ]
                return np.random.choice(responses)

        # ===== ABOUT/HELP QUESTIONS =====
        about_patterns = [
            ['what are you', 'who are you', 'tell me about yourself', 'what do you do', 'what can you help with', 'what can you do', 'how can you help', 'what are your capabilities', 'help', 'help me', 'i need help', 'can you help', 'how do i use this', 'how does this work', 'what should i ask']
        ]

        for pattern_group in about_patterns:
            if any(keyword in query_lower for keyword in pattern_group):
                return """🤖 **About BioGPT Medical Assistant**

I'm an AI medical assistant powered by BioGPT-Large, a specialized medical AI model trained on extensive medical literature. Here's what I can help you with:

🩺 **Medical Specialties:**
• Pediatric medicine and children's health
• Symptom explanation and medical conditions
• Treatment options and medical procedures
• When to seek medical care
• Prevention and wellness guidance

🎯 **How to Use Me:**
• Ask specific medical questions: "What causes fever in children?"
• Describe symptoms: "My child has a persistent cough"
• Seek guidance: "When should I call the doctor?"
• Get information: "How do I treat dehydration?"

⚠️ **Important Reminder:**
I provide educational medical information based on medical literature, but I'm not a substitute for professional medical advice. Always consult qualified healthcare providers for:
• Medical emergencies
• Diagnosis and treatment decisions
• Personalized medical advice
• Medication guidance

💡 **Tip:** Be specific in your questions for the most helpful responses!

What pediatric health topic would you like to explore?"""

        # ===== SMALL TALK & PERSONAL QUESTIONS =====
        personal_patterns = [
            ['how are you feeling', 'are you okay', 'how\'s your day', 'are you smart', 'are you intelligent', 'do you know everything', 'are you human', 'are you real', 'are you a robot', 'are you ai', 'you\'re smart', 'you\'re helpful', 'good job', 'well done', 'impressive']
        ]

        for pattern_group in personal_patterns:
            if any(keyword in query_lower for keyword in pattern_group):
                responses = [
                    "🤖 I'm an AI medical assistant, so I don't have feelings, but I'm functioning well and ready to help with medical questions! My purpose is to provide reliable pediatric health information. What can I help you with?",
                    "😊 Thank you for asking! As an AI, I'm always ready to assist with medical information. I'm designed to help with pediatric health questions using evidence-based medical knowledge. How can I help you today?",
                    "🎯 I'm doing what I do best - providing medical information! I'm an AI trained on medical literature to help with pediatric health questions. What medical topic interests you?"
                ]
                return np.random.choice(responses)

        # ===== CONFUSED/UNCLEAR INPUT =====
        confusion_patterns = [
            ['i don\'t know', 'not sure', 'confused', 'unclear', 'help me understand', 'what do you mean', 'i don\'t understand', 'can you explain', 'huh', 'i\'m lost', 'i\'m confused', 'this is confusing']
        ]

        for pattern_group in confusion_patterns:
            if any(keyword in query_lower for keyword in pattern_group):
                return """🤔 **I understand it can be confusing!** Let me help you get started.

💡 **Try asking questions like:**

🩺 **Symptoms:**
• "What causes [symptom] in children?"
• "My child has [symptom], what should I do?"

💊 **Treatments:**
• "How do I treat [condition] in children?"
• "What are treatment options for [condition]?"

🚨 **Urgency:**
• "When should I call the doctor about [symptom]?"
• "Is [symptom] serious in children?"

🛡️ **Prevention:**
• "How can I prevent [condition]?"
• "What are the warning signs of [condition]?"

**What specific aspect of your child's health would you like to understand better?**"""

        # ===== APOLOGIES & POLITENESS =====
        polite_patterns = [
            ['sorry', 'excuse me', 'pardon me', 'my apologies', 'please help', 'could you please', 'would you mind', 'if you don\'t mind', 'sorry to bother you']
        ]

        for pattern_group in polite_patterns:
            if any(keyword in query_lower for keyword in pattern_group):
                return "😊 No need to apologize! I'm here to help with medical questions. Please feel free to ask anything about pediatric health - that's exactly what I'm designed for. What can I help you with?"

        # ===== TESTING & VERIFICATION =====
        test_patterns = [
            ['test', 'testing', 'hello world', 'can you hear me', 'are you working', 'do you work', 'are you there', 'are you online', 'check', 'verify', 'ping']
        ]

        for pattern_group in test_patterns:
            if any(keyword in query_lower for keyword in pattern_group):
                return "✅ **System Check:** I'm working perfectly and ready to assist! BioGPT medical AI is online and functioning optimally. Ready to help with pediatric medical questions. What would you like to know?"

        # Return None if no conversational pattern matches
        return None


    def chat(self, query: str) -> str:
        """Main chat function with BioGPT and comprehensive conversational handling"""
        if not query.strip():
            return "Hello! I'm BioGPT, your professional medical AI assistant. How can I help you with pediatric medical questions today?"

        # Handle comprehensive conversational interactions first
        conversational_response = self.handle_conversational_interactions(query)
        if conversational_response:
            # Add to conversation history
            self.conversation_history.append({
                'query': query,
                'response': conversational_response,
                'timestamp': datetime.now().isoformat(),
                'type': 'conversational'
            })
            return conversational_response

        if not self.knowledge_chunks:
            return "Please load medical data first to access the medical knowledge base."

        if not self.model or not self.tokenizer:
            return "Medical model not available. Please check the setup."

        print(f"🔍 Processing medical query: {query}")

        # Retrieve relevant medical context using FAISS or keyword search
        start_time = time.time()
        context = self.retrieve_medical_context(query)
        retrieval_time = time.time() - start_time

        if not context:
            return "I don't have specific information about this topic in my medical database. Please consult with a healthcare professional for personalized medical advice."

        print(f"   📚 Context retrieved ({retrieval_time:.2f}s)")

        # Generate response with BioGPT
        start_time = time.time()
        main_context = '\n\n'.join(context)
        response = self.generate_biogpt_response(main_context, query)
        generation_time = time.time() - start_time

        print(f"   🧠 Response generated ({generation_time:.2f}s)")

        # Format final response
        final_response = f"🩺 **Medical Information:** {response}\n\n⚠️ **Important:** This information is for educational purposes only. Always consult with qualified healthcare professionals for medical diagnosis, treatment, and personalized advice."

        # Add to conversation history
        self.conversation_history.append({
            'query': query,
            'response': final_response,
            'timestamp': datetime.now().isoformat(),
            'retrieval_time': retrieval_time,
            'generation_time': generation_time,
            'type': 'medical'
        })

        return final_response

    def get_conversation_summary(self) -> Dict:
        """Get conversation statistics"""
        if not self.conversation_history:
            return {"message": "No conversations yet"}

        # Filter medical conversations for performance stats
        medical_conversations = [h for h in self.conversation_history if h.get('type') == 'medical']

        model_info = "BioGPT-Large" if self.model else "Model Not Loaded"

        if not medical_conversations:
            return {
                "total_conversations": len(self.conversation_history),
                "medical_conversations": 0,
                "conversational_interactions": len(self.conversation_history),
                "model_info": model_info,
                "vector_search": "FAISS CPU" if self.faiss_ready else "Keyword Search",
                "device": self.device
            }

        avg_retrieval_time = sum(h.get('retrieval_time', 0) for h in medical_conversations) / len(medical_conversations)
        avg_generation_time = sum(h.get('generation_time', 0) for h in medical_conversations) / len(medical_conversations)

        return {
            "total_conversations": len(self.conversation_history),
            "medical_conversations": len(medical_conversations),
            "conversational_interactions": len(self.conversation_history) - len(medical_conversations),
            "avg_retrieval_time": f"{avg_retrieval_time:.2f}s",
            "avg_generation_time": f"{avg_generation_time:.2f}s",
            "model_info": model_info,
            "vector_search": "FAISS CPU" if self.faiss_ready else "Keyword Search",
            "device": self.device,
            "quantization": "8-bit" if self.use_8bit else "16-bit/32-bit"
        }

# Create and Test BioGPT Chatbot

def create_biogpt_chatbot():
    """Create and initialize the BioGPT chatbot"""
    print("🚀 Creating Professional BioGPT Medical Chatbot")
    print("=" * 60)

    # Create chatbot
    chatbot = ColabBioGPTChatbot(use_gpu=True, use_8bit=True)

    return chatbot

def test_biogpt_chatbot(chatbot, test_file='Pediatric_cleaned.txt'):
    """Test the BioGPT chatbot"""
    print("\n📚 Loading medical data...")
    success = chatbot.load_medical_data(test_file)

    if not success:
        print("❌ Failed to load medical data. Please check the file.")
        return None

    print("\n🧪 Testing BioGPT Medical Chatbot:")
    print("=" * 50)

    # Test queries
    test_queries = [
        "What causes fever in children?",
        "How should I treat my child's cough?",
        "When should I be concerned about my baby's breathing?",
        "What are the signs of dehydration in infants?"
    ]

    for i, query in enumerate(test_queries, 1):
        print(f"\n{i}️⃣ Testing: {query}")
        print("-" * 40)

        response = chatbot.chat(query)
        print(f"🤖 BioGPT Response:\n{response}")
        print("=" * 50)

    # Show conversation summary
    summary = chatbot.get_conversation_summary()
    print("\n📊 Performance Summary:")
    for key, value in summary.items():
        print(f"   {key}: {value}")

    return chatbot

# Interactive Chat Interface

def interactive_biogpt_chat(chatbot):
    """Interactive chat with BioGPT"""
    print("\n💬 Interactive BioGPT Medical Chat")
    print("=" * 50)
    print("You're now chatting with BioGPT, a professional medical AI!")
    print("Type 'quit' to exit, 'summary' to see stats")
    print("-" * 50)

    while True:
        user_input = input("\n👤 You: ").strip()

        if user_input.lower() in ['quit', 'exit', 'bye']:
            print("\n👋 Thank you for using BioGPT Medical Assistant!")
            # Show final summary
            summary = chatbot.get_conversation_summary()
            print("\n📊 Final Session Summary:")
            for key, value in summary.items():
                print(f"   {key}: {value}")
            break

        elif user_input.lower() == 'summary':
            summary = chatbot.get_conversation_summary()
            print("\n📊 Current Session Summary:")
            for key, value in summary.items():
                print(f"   {key}: {value}")
            continue

        elif not user_input:
            continue

        print(f"\n🤖 BioGPT: ", end="")
        response = chatbot.chat(user_input)
        print(response)

# Main Execution - Initialization

import torch

# Create the BioGPT chatbot
chatbot = create_biogpt_chatbot()

print("\n" + "="*60)
print("🎯 NEXT STEPS:")
print("1. Upload your medical data file by running the next cell.")
print("2. Test the chatbot by running the cell after the upload.")
print("3. Start interactive chat by running the final cell.")
print("="*60)