Update app.py

app.py

@@ -1,164 +1,271 @@
-# Standard Libraries
+# Standard imports first
 import os
-import json
-import time
-import asyncio
+import torch
 import logging
-import gc
-import re
-import traceback
-from pathlib import Path
 from datetime import datetime
-from typing import List, Dict, Union, Tuple, Optional, Any
-from dataclasses import dataclass, field
-import zipfile
-
-# Machine Learning and Deep Learning Libraries
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from torch.cuda.amp import autocast
-from torch.utils.data import DataLoader
-
-# Hugging Face and Transformers
-from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
-from sentence_transformers import SentenceTransformer
-from datasets import load_dataset, Dataset, concatenate_datasets
 from huggingface_hub import login
-
-# FAISS and PEFT
-import faiss
-from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, TaskType, PeftModel
-
-# LangChain - updated imports as per recent deprecations
-from langchain_community.vectorstores import FAISS  # Updated import
-from langchain_community.embeddings import HuggingFaceEmbeddings  # Updated import
-from langchain_community.document_loaders import TextLoader  # Updated import
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-
-
-# External Tools and APIs
-import wandb
-import requests
-import gradio as gr
-import IPython.display as display  # Required for IPython display functionality
 from dotenv import load_dotenv
+from datasets import load_dataset, Dataset
+from transformers import (
+    AutoTokenizer, 
+    AutoModelForCausalLM, 
+    TrainingArguments, 
+    Trainer,
+    BitsAndBytesConfig
+)
+from peft import (
+    LoraConfig,
+    get_peft_model,
+    prepare_model_for_kbit_training
+)
 from tqdm.auto import tqdm
 
-# Suppress Warnings
-import warnings
-warnings.filterwarnings('ignore')
-
-
-# Ensure Hugging Face login
-try:
-    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)
-
-
-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'
-
-
-
 # Setup logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-
-
-class ModelManager:
-    """Handles model loading and resource management"""
-    
-    @staticmethod
-    def verify_and_extract_model(checkpoint_zip_path: str, extracted_model_dir: str) -> str:
-        """Verify and extract the model if it's not already extracted"""
-        if not os.path.exists(extracted_model_dir):
-            # Unzip the model if it hasn’t been extracted yet
-            with zipfile.ZipFile(checkpoint_zip_path, 'r') as zip_ref:
-                zip_ref.extractall(extracted_model_dir)
-            logger.info(f"Extracted model to: {extracted_model_dir}")
-        else:
-            logger.info(f"Model already extracted: {extracted_model_dir}")
-            
-        return extracted_model_dir
+class SecretsManager:
+    """Handles authentication and secrets management"""
     
     @staticmethod
-    def clear_gpu_memory():
-        """Clear GPU memory cache"""
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-            gc.collect()
-
-class PearlyBot:
-    def __init__(self):
-        try:
-            # Use the correct model path from your space
-            self.repo_id = "Pearilsa/pearly_med_triage_chatbot_kagglex"
-            self.model_filename = "pearly_model.zip"
-            self.setup_model()
-            self.setup_rag()
-            self.conversation_history = []
-            self.last_interaction_time = time.time()
-            self.interaction_cooldown = 1.0
-        except Exception as e:
-            logger.error(f"Error initializing bot: {e}")
-            raise
-
-    def setup_model(self):
-        """Initialize model from Hugging Face space"""
+    def setup_credentials():
+        """Setup all required credentials"""
         try:
-            logger.info(f"Loading model from {self.repo_id}")
+            # Load environment variables
+            load_dotenv()
             
-            # Download and prepare model path
-            local_model_path = os.path.join(os.getcwd(), "models")
-            os.makedirs(local_model_path, exist_ok=True)
+            # Get credentials
+            credentials = {
+                'KAGGLE_USERNAME': os.getenv('KAGGLE_USERNAME'),
+                'KAGGLE_KEY': os.getenv('KAGGLE_KEY'),
+                'HF_TOKEN': os.getenv('HF_TOKEN'),
+                'WANDB_KEY': os.getenv('WANDB_KEY')
+            }
             
-            # Load tokenizer and model from the space
-            self.tokenizer = AutoTokenizer.from_pretrained(
-                self.repo_id,
-                token=os.getenv("HF_TOKEN"),  # Use your Hugging Face token
-                cache_dir=local_model_path
-            )
-            self.tokenizer.pad_token = self.tokenizer.eos_token
-            logger.info("Tokenizer loaded successfully")
-            
-            # Load model with 8-bit quantization
-            self.model = AutoModelForCausalLM.from_pretrained(
-                self.repo_id,
-                token=os.getenv("HF_TOKEN"),
-                device_map="auto",
-                load_in_8bit=True,
-                torch_dtype=torch.float16,
-                low_cpu_mem_usage=True,
-                cache_dir=local_model_path
-            )
-            self.model.eval()
-            logger.info("Model loaded successfully")
+            # Validate credentials
+            missing_creds = [k for k, v in credentials.items() if not v]
+            if missing_creds:
+                logger.warning(f"Missing credentials: {', '.join(missing_creds)}")
+            
+            # Setup Hugging Face authentication
+            if credentials['HF_TOKEN']:
+                login(token=credentials['HF_TOKEN'])
+                logger.info("Successfully logged in to Hugging Face")
+# Setup Kaggle credentials if available
+            if credentials['KAGGLE_USERNAME'] and credentials['KAGGLE_KEY']:
+                os.environ['KAGGLE_USERNAME'] = credentials['KAGGLE_USERNAME']
+                os.environ['KAGGLE_KEY'] = credentials['KAGGLE_KEY']
+            
+            # Setup wandb if available
+            if credentials['WANDB_KEY']:
+                os.environ['WANDB_API_KEY'] = credentials['WANDB_KEY']
+            
+            return credentials
             
         except Exception as e:
-            logger.error(f"Error in model setup: {str(e)}")
+            logger.error(f"Error setting up credentials: {e}")
             raise
+class ModelTrainer:
+    """Handles model training pipeline"""
+    
+    def __init__(self):
+        # Set memory optimization environment variables
+        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'
+        
+        # Initialize attributes
+        self.model = None
+        self.tokenizer = None
+        self.dataset = None
+        self.processed_dataset = None
+        self.chunk_size = 300
+        self.chunk_overlap = 100
+        self.num_relevant_chunks = 3
+        self.vector_store = None
+        self.embeddings = None
+        self.last_interaction_time = time.time()  # Add this
+        self.interaction_cooldown = 1.0  # Add this
+        
+        # Setup GPU preferences
+        torch.backends.cuda.matmul.allow_tf32 = False
+        torch.backends.cudnn.allow_tf32 = False
+
+    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
             
+        return results
+    try:
+         Load and process each dataset separately
+    print("Processing MedQA dataset...")
+    medqa = load_dataset("medalpaca/medical_meadow_medqa", split="train[:500]")
+    medqa_processed = []
+    
+    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()
+    
+    print("Processing MedDiagnosis dataset...")
+    meddia = load_dataset("wasiqnauman/medical-diagnosis-synthetic", split="train[:500]")
+    meddia_processed = []
+    
+    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()
+    
+    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]
+                
+                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
+        
     def setup_rag(self):
+        """Initialize RAG components"""
         try:
-            # Add configuration options
-            self.chunk_size = 300
-            self.chunk_overlap = 100
-            self.num_relevant_chunks = 3
+            logger.info("Setting up RAG system...")
             
             # Load knowledge base
             knowledge_base = self._load_knowledge_base()
             
-            # Setup embeddings with error handling
+            # Setup embeddings
             self.embeddings = self._initialize_embeddings()
             
-            # Enhanced text splitting
+            # Process texts for vector store
             texts = self._split_texts(knowledge_base)
             
             # Create vector store with metadata
@@ -168,13 +275,15 @@ class PearlyBot:
                 metadatas=[{"source": f"chunk_{i}"} for i in range(len(texts))]
             )
             
-            # Add validation
+            # Validate RAG setup
             self._validate_rag_setup()
+            logger.info("RAG system setup complete")
             
         except Exception as e:
-            logger.error(f"RAG setup failed: {str(e)}")
+            logger.error(f"Failed to setup RAG: {e}")
             raise
-            # Load your knowledge base content
+
+           # Load your knowledge base content
     def _load_knowledge_base(self):
         """Load and validate knowledge base content"""
         try:
@@ -488,6 +597,156 @@ class PearlyBot:
         except Exception as e:
             logger.error(f"RAG system validation failed: {str(e)}")
             raise
+    
+
+
+    
+
+    
+    
+    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 train(self):
+        """Main training pipeline with RAG integration"""
+        try:
+            logger.info("Starting training pipeline")
+            
+            # Clear GPU memory
+            torch.cuda.empty_cache()
+            if torch.cuda.is_available():
+                torch.cuda.reset_peak_memory_stats()
+            
+            # Setup model, tokenizer, and RAG
+            logger.info("Setting up model components...")
+            self.model, self.tokenizer = self.setup_model_and_tokenizer()
+            self.setup_rag()
+            
+            # Prepare and process datasets
+            logger.info("Preparing datasets...")
+            self.dataset = self.prepare_initial_datasets(batch_size=4)
+            self.processed_dataset = self.prepare_dataset(
+                self.dataset,
+                self.tokenizer,
+                max_length=256,
+                batch_size=2
+            )
+            
+            # Train model
+            logger.info("Starting training...")
+            training_args = self.setup_training_arguments()
+            trainer = Trainer(
+                model=self.model,
+                args=training_args,
+                train_dataset=self.processed_dataset,
+                tokenizer=self.tokenizer
+            )
+            trainer.train()
+            
+            # Save and push to hub
+            logger.info("Saving model...")
+            trainer.save_model()
+            if os.getenv('HF_TOKEN'):
+                trainer.push_to_hub(
+                    "Pearilsa/pearly_med_triage_chatbot_kagglex",
+                    private=True
+                )
+            
+            logger.info("Training completed successfully!")
+            
+        except Exception as e:
+            logger.error(f"Training failed: {e}")
+            raise
+        finally:
+            torch.cuda.empty_cache()
+
+if __name__ == "__main__":
+    # Initialize trainer
+    trainer = ModelTrainer()
+    
+    # Train model
+    trainer.train()
+    
+    def _get_enhanced_context(self, query: str) -> str:
+        """Get relevant context with scores"""
+        try:
+            # Get documents with similarity scores
+            docs_and_scores = self.vector_store.similarity_search_with_score(
+                query,
+                k=self.num_relevant_chunks
+            )
+            
+            # Filter and format relevant contexts
+            relevant_contexts = []
+            for doc, score in docs_and_scores:
+                if score < 0.8:  # Lower score means more relevant
+                    source = doc.metadata.get('source', 'Unknown')
+                    relevant_contexts.append(
+                        f"[Source: {source}]\n{doc.page_content}"
+                    )
+            
+            return "\n\n".join(relevant_contexts) if relevant_contexts else ""
+            
+        except Exception as e:
+            logger.error(f"Error retrieving enhanced context: {e}")
+            return ""
 
     def _initialize_embeddings(self):
         try:
@@ -525,24 +784,22 @@ class PearlyBot:
     def generate_response(self, message: str, history: list) -> str:
         """Generate response using both fine-tuned model and RAG"""
         try:
-            # Rate limiting
+            # Rate limiting and memory management
             current_time = time.time()
             if current_time - self.last_interaction_time < self.interaction_cooldown:
                 time.sleep(self.interaction_cooldown)
+            torch.cuda.empty_cache()
             
-            # Clear GPU memory before generation
-            ModelManager.clear_gpu_memory()
-            
-            # Get RAG context
-            context = self.get_relevant_context(message)
+            # Get enhanced context from RAG
+            context = self._get_enhanced_context(message)
             
             # Format conversation history
             conv_history = "\n".join([
-                f"User: {user}\nAssistant: {assistant}"
-                for user, assistant in history[-3:]  # Keep last 3 turns
+                f"User: {turn['input']}\nAssistant: {turn['output']}"
+                for turn in history[-3:]  # Keep last 3 turns
             ])
             
-            # Create prompt
+            # Create enhanced prompt with RAG context
             prompt = f"""<start_of_turn>system
 Using these medical guidelines:
 
@@ -552,9 +809,9 @@ Previous conversation:
 {conv_history}
 
 Guidelines:
-1. Assess symptoms and severity
-2. Ask relevant follow-up questions
-3. Direct to appropriate care (999, 111, or GP)
+1. Assess symptoms and severity based on both your training and the provided guidelines
+2. Ask relevant follow-up questions if needed
+3. Direct to appropriate care (999, 111, or GP) according to symptom severity
 4. Show empathy and cultural sensitivity
 5. Never diagnose or recommend treatments
 <end_of_turn>
@@ -563,41 +820,36 @@ Guidelines:
 <end_of_turn>
 <start_of_turn>assistant"""
 
-            # Generate response
-            try:
-                inputs = self.tokenizer(
-                    prompt,
-                    return_tensors="pt",
-                    truncation=True,
-                    max_length=512
-                ).to(self.model.device)
-
-                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,
-                    no_repeat_ngram_size=3
-                )
-                
-                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.last_interaction_time = time.time()
-                return response
-                
-            except torch.cuda.OutOfMemoryError:
-                ModelManager.clear_gpu_memory()
-                logger.error("GPU out of memory, cleared cache and retrying...")
-                return "I apologize, but I'm experiencing technical difficulties. Please try again."
-                
+            # Generate response with model
+            inputs = self.tokenizer(
+                prompt,
+                return_tensors="pt",
+                truncation=True,
+                max_length=512
+            ).to(self.model.device)
+
+            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,
+                no_repeat_ngram_size=3
+            )
+            
+            # Process response
+            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.last_interaction_time = time.time()
+            return response
+            
         except Exception as e:
-            logger.error(f"Error generating response: {str(e)}")
+            logger.error(f"Error generating response: {e}")
             return "I apologize, but I encountered an error. Please try again."
             
     def handle_feedback(self, message: str, response: str, feedback: int):
@@ -971,23 +1223,10 @@ def create_demo():
         raise
 
 if __name__ == "__main__":
-    try:
-        # Initialize logging
-        logging.basicConfig(level=logging.INFO)
-        
-        # Load environment variables
-        load_dotenv()
-        
-        # Create and launch demo
-        demo = create_demo()
-        demo.launch(
-            server_name="0.0.0.0",
-            server_port=7860,
-            show_error=True
-        )
-        
-    except Exception as e:
-        logger.error(f"Application startup failed: {e}")
-        raise
-
-
+    # Initialize logging and load env vars
+    logging.basicConfig(level=logging.INFO)
+    load_dotenv()
+    
+    # Create and launch demo
+    demo = create_demo()
+    demo.launch(server_name="0.0.0.0", server_port=7860)