Update app.py

app.py

@@ -1,251 +1,472 @@
 # app.py
 import os
-import logging
+import json
+import keras
+from datasets import load_dataset
+import tensorflow as tf
+from huggingface_hub import login
 import torch
-from typing import Dict, List, Any
-import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from transformers import ( AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer)
 from sentence_transformers import SentenceTransformer
-from peft import get_peft_model, LoraConfig, TaskType, prepare_model_for_kbit_training
+from typing import List, Dict, Union, Tuple
 import faiss
 import numpy as np
-from datasets import load_dataset
+from datasets import Dataset
+import torch.nn.functional as F
+from torch.cuda.amp import autocast
+import gc
+from peft import ( LoraConfig, get_peft_model, prepare_model_for_kbit_training, TaskType, PeftModel)
+from tqdm.auto import tqdm
+from torch.utils.data import DataLoader
+import logging
+import wandb
+from pathlib import Path
+from typing import List, Dict, Union, Optional, Any
+import torch.nn as nn
+from dataclasses import dataclass, field
+import time
+import asyncio
+import pytest
+from unittest.mock import Mock, patch
+from sklearn.metrics import classification_report, confusion_matrix
+import gradio as gr
+import matplotlib.pyplot as plt
 from datetime import datetime
-import json
-from huggingface_hub import login
-from dotenv import load_dotenv
-
-# Quick connectivity test to Hugging Face
+import requests
+import pandas as pd
+import seaborn as sns
+import traceback
+from matplotlib.gridspec import GridSpec
+from datasets import load_dataset, concatenate_datasets
+from langchain.vectorstores import FAISS
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.document_loaders import TextLoader
+from google.colab import output
+import IPython.display as display
+from peft import get_peft_model, LoraConfig, prepare_model_for_kbit_training
+
+
+# Ensure Hugging Face login
 try:
-    import huggingface_hub
-    whoami = huggingface_hub.whoami()
-    print("Hugging Face Login successful:", whoami)
+    hf_token = os.getenv("HF_TOKEN")
+    if hf_token:
+        login(token=hf_token)
+    print("Login successful!")
 except Exception as e:
     print("Hugging Face Login failed:", e)
 
+# CUDA and Memory Configurations
+torch.backends.cuda.matmul.allow_tf32 = False
+torch.backends.cudnn.allow_tf32 = False
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:64,garbage_collection_threshold:0.8,expandable_segments:True'
+os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
+
 
-# Load environment variables
-load_dotenv()
-
-# Suppress TensorFlow INFO and WARNING logs
-os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
-os.environ.pop("HF_HUB_OFFLINE", None)
-
-# Configure logging
-logging.basicConfig(
-    level=logging.INFO,
-    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
-    handlers=[
-        logging.StreamHandler(),
-        logging.FileHandler('app.log')
-    ]
-)
-
-# Create a logger for your application
-logger = logging.getLogger('medical_bot')
-
-# Retrieve secrets securely from environment variables
-hf_token = os.getenv("HF_TOKEN")
-if hf_token:
-    login(token=hf_token)
-
-
-
-class AdaptiveMedicalBot:
-    def __init__(self):
-        self.config = self.AdaptiveBotConfig()
-        self.setup_models()
-        self.load_datasets()
-        self.setup_adaptive_learning()
-        self.conversation_history = []  # Store conversation history
-        self.symptom_tracker = {}  # Track symptoms and severity across conversation
-
-    class AdaptiveBotConfig:
-        MODEL_NAME = "google/gemma-7b"
-        EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
-        LORA_R = 8
-        LORA_ALPHA = 16
-        LORA_DROPOUT = 0.1
-        LORA_TARGET_MODULES = ["q_proj", "v_proj"]
-        MAX_LENGTH = 512
-        BATCH_SIZE = 1
-        LEARNING_RATE = 1e-4
-
-    def setup_adaptive_learning(self):
-        """Initialize adaptive learning components"""
-        self.feedback_history = []
-
-    def setup_models(self):
-        """Initialize models with LoRA and quantization"""
-        try:
-            # Add explicit device check
-            device = "cuda" if torch.cuda.is_available() else "cpu"
-            logger.info(f"Using device: {device}")
+def prepare_initial_datasets(batch_size=8):
+    print("Loading datasets with memory-optimized batch processing...")
+    
+    def process_medqa_batch(examples):
+        results = []
+        inputs = examples['input']
+        instructions = examples['instruction']
+        outputs = examples['output']
+        
+        for inp, inst, out in zip(inputs, instructions, outputs):
+            results.append({
+                "input": f"{inp} {inst}",
+                "output": out
+            })
+        return results
+
+    def process_meddia_batch(examples):
+        results = []
+        inputs = examples['input']
+        outputs = examples['output']
+        
+        for inp, out in zip(inputs, outputs):
+            results.append({
+                "input": inp,
+                "output": out
+            })
+        return results
+
+    def process_persona_batch(examples):
+        results = []
+        personalities = examples['personality']
+        utterances = examples['utterances']
+        
+        for pers, utts in zip(personalities, utterances):
+            try:
+                # Process personality list
+                personality = ' '.join([
+                    p for p in pers 
+                    if isinstance(p, str)
+                ])
+                
+                # Process utterances
+                if utts and len(utts) > 0:
+                    utterance = utts[0]
+                    history = []
+                    
+                    # Process history
+                    if 'history' in utterance and utterance['history']:
+                        history = [
+                            h for h in utterance['history']
+                            if isinstance(h, str)
+                        ]
+                    
+                    history_text = ' '.join(history)
+                    
+                    # Get candidate response
+                    candidate = utterance.get('candidates', [''])[0] if utterance.get('candidates') else ''
+                    
+                    if personality or history_text:
+                        results.append({
+                            "input": f"{personality} {history_text}".strip(),
+                            "output": candidate
+                        })
+            except Exception as e:
+                print(f"Error processing persona batch item: {e}")
+                continue
             
-            bnb_config = BitsAndBytesConfig(
-                load_in_4bit=True,
-                bnb_4bit_quant_type="nf4",
-                bnb_4bit_compute_dtype=torch.float16,
-                bnb_4bit_use_double_quant=True
-            )
+        return results
+
+    # Load and process each dataset separately
+    print("Processing MedQA dataset...")
+    medqa = load_dataset("medalpaca/medical_meadow_medqa", split="train[:500]")
+    medqa_processed = []
     
-            # Add loading message
-            logger.info("Loading tokenizer...")
-            self.tokenizer = AutoTokenizer.from_pretrained(
-                self.config.MODEL_NAME,
-                use_auth_token=True,
-                trust_remote_code=True
-            )
+    for i in tqdm(range(0, len(medqa), batch_size), desc="Processing MedQA"):
+        batch = medqa[i:i + batch_size]
+        medqa_processed.extend(process_medqa_batch(batch))
+        if i % (batch_size * 5) == 0:
+            torch.cuda.empty_cache()
     
-            logger.info("Loading base model...")
-            base_model = AutoModelForCausalLM.from_pretrained(
-                self.config.MODEL_NAME,
-                quantization_config=bnb_config,
-                device_map="auto",
-                trust_remote_code=True,
-                use_auth_token=True
-            )
+    print("Processing MedDiagnosis dataset...")
+    meddia = load_dataset("wasiqnauman/medical-diagnosis-synthetic", split="train[:500]")
+    meddia_processed = []
     
-            logger.info("Preparing model for k-bit training...")
-            base_model = prepare_model_for_kbit_training(base_model)
-            
-            logger.info("Configuring LoRA...")
-            lora_config = LoraConfig(
-                r=self.config.LORA_R,
-                lora_alpha=self.config.LORA_ALPHA,
-                target_modules=self.config.LORA_TARGET_MODULES,
-                lora_dropout=self.config.LORA_DROPOUT,
-                bias="none",
-                task_type=TaskType.CAUSAL_LM
-            )
+    for i in tqdm(range(0, len(meddia), batch_size), desc="Processing MedDiagnosis"):
+        batch = meddia[i:i + batch_size]
+        meddia_processed.extend(process_meddia_batch(batch))
+        if i % (batch_size * 5) == 0:
+            torch.cuda.empty_cache()
     
-            logger.info("Applying PEFT configuration...")
-            self.model = get_peft_model(base_model, lora_config)
-            
-            logger.info("Loading embedding model...")
-            self.embedding_model = SentenceTransformer(
-                self.config.EMBEDDING_MODEL,
-                device=device
-            )
-            
-            logger.info("Models loaded successfully!")
-            
-        except Exception as e:
-            logger.error(f"Error setting up models: {e}")
-            raise
-
-    def load_datasets(self):
-        """Load and prepare datasets for RAG"""
-        try:
-            datasets = {
-                "medqa": load_dataset("medalpaca/medical_meadow_medqa", split="train[:500]"),
-                "diagnosis": load_dataset("wasiqnauman/medical-diagnosis-synthetic", split="train[:500]"),
-                "persona": load_dataset("AlekseyKorshuk/persona-chat", split="train[:500]")
-            }
-
-            self.documents = []
-            for dataset_name, dataset in datasets.items():
-                for item in dataset:
-                    if dataset_name == "persona":
-                        if isinstance(item.get('personality'), list):
-                            self.documents.append({'text': " ".join(item['personality']), 'type': 'persona'})
-                    else:
-                        if 'input' in item and 'output' in item:
-                            self.documents.append({'text': f"{item['input']}\n{item['output']}", 'type': dataset_name})
-
-            self._create_index()
-        except Exception as e:
-            logger.error(f"Error loading datasets: {e}")
-            raise
-
-    def _create_index(self):
-        """Create FAISS index for RAG"""
-        try:
-            sample_embedding = self.embedding_model.encode("sample text")
-            self.index = faiss.IndexFlatIP(sample_embedding.shape[0])
-
-            embeddings = [self.embedding_model.encode(doc['text']) for doc in self.documents]
-            self.index.add(np.array(embeddings))
-        except Exception as e:
-            logger.error(f"Error creating FAISS index: {e}")
-            raise
-
-    def generate_follow_up_questions(self, message: str, context: Dict[str, Any]) -> List[str]:
-        """Generate follow-up questions based on context"""
-        try:
-            prompt = f"""Patient message: "{message}"
-            Generate relevant follow-up questions focusing on timing, severity, associated symptoms, and impact on daily life.
-            Questions:"""
-
-            inputs = self.tokenizer(prompt, return_tensors="pt", max_length=self.config.MAX_LENGTH).to(self.model.device)
-            outputs = self.model.generate(inputs['input_ids'], max_new_tokens=50, temperature=0.7, do_sample=True)
-            questions = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-            return questions.split("\n")
-        except Exception as e:
-            logger.error(f"Error generating follow-up questions: {e}")
-            return ["Could you tell me more about when this started?"]
-
-    def assess_symptom_severity(self, message: str) -> str:
-        """Assess severity based on keywords in the message"""
-        if "severe" in message.lower() or "emergency" in message.lower():
-            return "emergency"
-        elif "persistent" in message.lower() or "moderate" in message.lower():
-            return "urgent"
-        return "routine"
-
-    def generate_response(self, message: str) -> Dict[str, Any]:
-        """Generate a response based on the message"""
-        try:
-            severity = self.assess_symptom_severity(message)
-            response = ""
-
-            # Retrieve relevant documents from FAISS
-            query_embedding = self.embedding_model.encode([message])
-            _, indices = self.index.search(query_embedding, k=5)
-            relevant_docs = [self.documents[idx]['text'] for idx in indices[0]]
-
-            prompt = f"""As a compassionate medical assistant, analyze the patient message: "{message}".
-            Consider relevant knowledge and the following documents:\n{relevant_docs}.
-            Respond with empathy, follow-up questions, and care guidance."""
-
-            inputs = self.tokenizer(prompt, return_tensors="pt", max_length=self.config.MAX_LENGTH).to(self.model.device)
-            outputs = self.model.generate(inputs['input_ids'], max_new_tokens=100, temperature=0.7, do_sample=True)
-            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-            follow_ups = self.generate_follow_up_questions(message, {})
-            response += f"\n{follow_ups[0]}"
+    print("Processing Persona-Chat dataset...")
+    persona = load_dataset("AlekseyKorshuk/persona-chat", split="train[:500]")
+    persona_processed = []
+    
+    for i in tqdm(range(0, len(persona), batch_size), desc="Processing Persona-Chat"):
+        batch = persona[i:i + batch_size]
+        persona_processed.extend(process_persona_batch(batch))
+        if i % (batch_size * 5) == 0:
+            torch.cuda.empty_cache()
+    
+    torch.cuda.empty_cache()
+    
+    print("Creating final dataset...")
+    all_processed = persona_processed + medqa_processed + meddia_processed
+    
+    valid_data = {
+        "input": [],
+        "output": []
+    }
+    
+    for item in all_processed:
+        if item["input"].strip() and item["output"].strip():
+            valid_data["input"].append(item["input"])
+            valid_data["output"].append(item["output"])
+    
+    final_dataset = Dataset.from_dict(valid_data)
+    
+    print(f"Final dataset size: {len(final_dataset)}")
+    return final_dataset
+
+def prepare_dataset(dataset, tokenizer, max_length=256, batch_size=4):
+    def tokenize_batch(examples):
+        formatted_texts = []
+        
+        for i in range(0, len(examples['input']), batch_size):
+            sub_batch_inputs = examples['input'][i:i + batch_size]
+            sub_batch_outputs = examples['output'][i:i + batch_size]
             
-            # Append response to conversation history
-            self.conversation_history.append((message, response))
-
-            # Add care level guidance
-            if severity == "emergency":
-                response += "\nThis seems urgent. Please call 999 immediately."
-            elif severity == "urgent":
-                response += "\nConsider calling NHS 111 for urgent assistance."
-
-            return {'response': response}
-        except Exception as e:
-            logger.error(f"Error generating response: {e}")
-            return {
-                'response': "I'm experiencing technical issues. If this is an emergency, please call 999 immediately.",
-            }
-
-    def handle_feedback(self, message: str, response: str, feedback: int):
-        """Update model based on feedback"""
-        try:
-            self.feedback_history.append({
-                'message': message,
-                'response': response,
-                'feedback': feedback,
-                'timestamp': datetime.now().isoformat()
-            })
+            for input_text, output_text in zip(sub_batch_inputs, sub_batch_outputs):
+                try:
+                    formatted_text = f"""<start_of_turn>user
+{input_text}
+<end_of_turn>
+<start_of_turn>assistant
+{output_text}
+<end_of_turn>"""
+                    formatted_texts.append(formatted_text)
+                except Exception as e:
+                    print(f"Error formatting text: {e}")
+                    continue
+        
+        tokenized = tokenizer(
+            formatted_texts,
+            padding="max_length",
+            truncation=True,
+            max_length=max_length,
+            return_tensors=None
+        )
+        
+        tokenized["labels"] = tokenized["input_ids"].copy()
+        return tokenized
+    
+    print(f"Tokenizing dataset in small batches (size={batch_size})...")
+    tokenized_dataset = dataset.map(
+        tokenize_batch,
+        batched=True,
+        batch_size=batch_size,
+        remove_columns=dataset.column_names,
+        desc="Tokenizing dataset",
+        load_from_cache_file=False
+    )
+    
+    return tokenized_dataset
 
-            if len(self.feedback_history) >= 10:
-                # Implement learning updates from feedback
-                self.feedback_history = []  # Reset history after learning update
-        except Exception as e:
-            logger.error(f"Error processing feedback: {e}")
+def setup_model_and_tokenizer(model_name="google/gemma-2b"):
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    tokenizer.pad_token = tokenizer.eos_token
+    
+    from transformers import BitsAndBytesConfig
+    
+    bnb_config = BitsAndBytesConfig(
+        load_in_8bit=True,
+        bnb_8bit_compute_dtype=torch.float16,
+        llm_int8_enable_fp32_cpu_offload=True
+    )
+    
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        device_map="auto",
+        quantization_config=bnb_config,
+        torch_dtype=torch.float16,
+        low_cpu_mem_usage=True
+    )
+    
+    model = prepare_model_for_kbit_training(model)
+    
+    lora_config = LoraConfig(
+        r=4,
+        lora_alpha=16,
+        target_modules=["q_proj", "v_proj"],
+        lora_dropout=0.05,
+        bias="none",
+        task_type="CAUSAL_LM"
+    )
+    
+    model = get_peft_model(model, lora_config)
+    model.print_trainable_parameters()
+    
+    return model, tokenizer
+
+def setup_training_arguments(output_dir="./pearly_fine_tuned"):
+    return TrainingArguments(
+        output_dir=output_dir,
+        num_train_epochs=1,
+        per_device_train_batch_size=1,
+        gradient_accumulation_steps=16,
+        warmup_steps=50,
+        logging_steps=10,
+        save_steps=200,
+        learning_rate=2e-4,
+        fp16=True,
+        gradient_checkpointing=True,
+        gradient_checkpointing_kwargs={"use_reentrant": False},
+        optim="adamw_8bit",
+        max_grad_norm=0.3,
+        weight_decay=0.001,
+        logging_dir="./logs",
+        save_total_limit=2,
+        remove_unused_columns=False,
+        dataloader_pin_memory=False,
+        max_steps=500,
+        report_to=["none"],
+    )
+
+def main():
+    torch.backends.cuda.matmul.allow_tf32 = False
+    torch.backends.cudnn.allow_tf32 = False
+    
+    torch.cuda.empty_cache()
+    if torch.cuda.is_available():
+        torch.cuda.reset_peak_memory_stats()
+    
+    print("Preparing initial datasets...")
+    combined_dataset = prepare_initial_datasets(batch_size=4)
+    
+    print(f"\nDataset size: {len(combined_dataset)}")
+    print(f"Column names: {combined_dataset.column_names}")
+    
+    if len(combined_dataset) > 0:
+        print("\nSample input-output pair:")
+        print(f"Input: {combined_dataset[0]['input'][:100]}...")
+        print(f"Output: {combined_dataset[0]['output'][:100]}...")
+    
+    print("\nSetting up model and tokenizer...")
+    model, tokenizer = setup_model_and_tokenizer()
+    
+    print("\nPreparing dataset for training...")
+    processed_dataset = prepare_dataset(
+        combined_dataset, 
+        tokenizer, 
+        max_length=256,
+        batch_size=2
+    )
+    
+    torch.cuda.empty_cache()
+    
+    training_args = setup_training_arguments()
+    
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=processed_dataset,
+        tokenizer=tokenizer,
+    )
+    
+    print("\nStarting training...")
+    try:
+        trainer.train()
+    except Exception as e:
+        print(f"Training error: {e}")
+        torch.cuda.empty_cache()
+        raise e
+    finally:
+        torch.cuda.empty_cache()
+        
+    print("\nSaving model...")
+    trainer.save_model()
+    print("Training completed!")
+
+DISCLAIMER = """
+IMPORTANT MEDICAL DISCLAIMER:
+Pearly is an AI medical triage assistant designed to help direct you to appropriate medical services.
+Pearly DOES NOT:
+- Make medical diagnoses
+- Prescribe medications
+- Provide specific treatment recommendations
+- Replace professional medical advice
+
+Always consult qualified healthcare professionals for medical advice and treatment.
+In case of emergency, call 999 immediately.
+"""
+
+class PearlyBot:
+    def __init__(self, model_path="./pearly_fine_tuned", embedding_model="sentence-transformers/all-MiniLM-L6-v2"):
+        print("Loading saved model...")
+        print(DISCLAIMER)
+        
+        # Clean memory
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        
+        # Load tokenizer and model directly from saved path
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_path,
+            torch_dtype=torch.float16,
+            low_cpu_mem_usage=True,
+            device_map="auto"
+        )
+        
+        self.model.eval()  # Set to evaluation mode
+        
+        # Initialize RAG components
+        self.embeddings = HuggingFaceEmbeddings(model_name=embedding_model)
+        self.vector_store = None
+        self.conversation_history = []
+
+    def initialize_rag(self, documents_path="./knowledge_base"):
+        """Initialize RAG system"""
+        print("Loading knowledge base...")
+        
+        text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=300,
+            chunk_overlap=100,
+            separators=["\n\n", "\n", ".", "!", "?", ":"]
+        )
+        
+        documents = []
+        for filename in os.listdir(documents_path):
+            if filename.endswith('.txt'):
+                loader = TextLoader(os.path.join(documents_path, filename))
+                documents.extend(loader.load())
+        
+        texts = text_splitter.split_documents(documents)
+        self.vector_store = FAISS.from_documents(texts, self.embeddings)
+        self.retriever = self.vector_store.as_retriever(
+            search_type="similarity",
+            search_kwargs={"k": 5}
+        )
+        print("Knowledge base loaded successfully!")
+
+    def get_relevant_context(self, user_input):
+        if not self.retriever:
+            return ""
+        docs = self.retriever.get_relevant_documents(user_input)
+        return "\n\n".join([doc.page_content for doc in docs])
+
+    def generate_response(self, user_input):
+        context = self.get_relevant_context(user_input)
+        history = "\n".join([
+            f"User: {turn['user']}\nAssistant: {turn['assistant']}\n"
+            for turn in self.conversation_history[-3:]
+        ])
+        
+        prompt = f"""<start_of_turn>system
+As Pearly, I use the following medical guidelines to help triage patients:
+
+{context}
+
+Previous Conversation:
+{history}
+
+Based on these guidelines, I will:
+1. Assess symptoms and severity
+2. Ask relevant follow-up questions
+3. Direct to appropriate care (999, 111, or GP)
+4. Show empathy and cultural sensitivity
+5. Never diagnose or recommend treatments
+<end_of_turn>
+<start_of_turn>user
+{user_input}
+<end_of_turn>
+<start_of_turn>assistant"""
+
+        inputs = self.tokenizer(
+            prompt,
+            return_tensors="pt",
+            truncation=True,
+            max_length=512
+        ).to(self.model.device)
+        
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **inputs,
+                max_new_tokens=256,
+                min_new_tokens=20,
+                do_sample=True,
+                temperature=0.7,
+                top_p=0.9,
+                repetition_penalty=1.2,
+                pad_token_id=self.tokenizer.pad_token_id
+            )
+        
+        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        response = response.split("<start_of_turn>assistant")[-1].strip()
+        if "<end_of_turn>" in response:
+            response = response.split("<end_of_turn>")[0].strip()
+        
+        self.conversation_history.append({
+            "user": user_input,
+            "assistant": response
+        })
+        
+        return response
 
 def create_demo():
     """Set up Gradio interface for the chatbot with enhanced styling and functionality."""
@@ -607,17 +828,8 @@ def create_demo():
         raise
 
 if __name__ == "__main__":
-    # Initialize environment
-    load_dotenv()
-    
-    # Set up Hugging Face login if token exists
-    hf_token = os.getenv("HF_TOKEN")
-    if hf_token:
-        login(token=hf_token)
-    
-    # Launch demo
-    os.environ.pop("HF_HUB_OFFLINE", None)  # Ensure online mode
-    demo = create_demo()
+    load_dotenv()  # Load environment variables
+    demo = create_demo()  # Launch the Gradio app
     demo.launch(share=True)