Create app.py

app.py

@@ -0,0 +1,891 @@
+import os
+import subprocess
+import sys
+import warnings
+import logging
+from typing import List, Dict, Any, Optional
+import tempfile
+import re
+import time
+import gc
+import spaces
+
+# Set up logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.FileHandler("debug.log"),
+        logging.StreamHandler()
+    ]
+)
+logger = logging.getLogger(__name__)
+
+# Suppress warnings
+warnings.filterwarnings("ignore")
+
+def install_package(package: str, version: Optional[str] = None) -> None:
+    """Install a Python package if not already installed"""
+    package_spec = f"{package}=={version}" if version else package
+    try:
+        subprocess.check_call([sys.executable, "-m", "pip", "install", "--no-cache-dir", package_spec])
+        print(f"Successfully installed {package_spec}")
+    except subprocess.CalledProcessError as e:
+        print(f"Failed to install {package_spec}: {e}")
+        raise
+
+# Required packages - install these before importing
+required_packages = {
+    "torch": None,
+    "gradio": "3.10.1",
+    "transformers": None,
+    "peft": None,
+    "bitsandbytes": None,
+    "PyPDF2": None,
+    "python-docx": None,
+    "accelerate": None,
+    "sentencepiece": None,
+}
+
+# Install required packages BEFORE importing them
+for package, version in required_packages.items():
+    try:
+        __import__(package)
+        print(f"{package} is already installed.")
+    except ImportError:
+        print(f"Installing {package}...")
+        install_package(package, version)
+
+# Now we can safely import all required modules
+import torch
+import transformers
+import gradio as gr
+from transformers import (
+    AutoTokenizer, AutoModelForCausalLM,
+    TrainingArguments, Trainer, TrainerCallback,
+    BitsAndBytesConfig
+)
+from peft import (
+    LoraConfig,
+    prepare_model_for_kbit_training,
+    get_peft_model
+)
+import PyPDF2
+import docx
+import numpy as np
+from tqdm import tqdm
+from torch.utils.data import Dataset as TorchDataset
+
+# Suppress transformers warnings
+transformers.logging.set_verbosity_error()
+
+# Check GPU availability
+if torch.cuda.is_available():
+    DEVICE = "cuda"
+    print(f"GPU found: {torch.cuda.get_device_name(0)}")
+    print(f"CUDA version: {torch.version.cuda}")
+else:
+    DEVICE = "cpu"
+    print("No GPU found, using CPU. Fine-tuning will be much slower.")
+    print("For better performance, use Google Colab with GPU runtime (Runtime > Change runtime type > GPU)")
+
+# Constants specific to Phi-2
+MODEL_KEY = "microsoft/phi-2"
+MAX_SEQ_LEN = 512  # Reduced from 1024 for much lighter memory usage
+# FIX: Updated target modules for Phi-2
+LORA_TARGET_MODULES = ["q_proj", "k_proj", "v_proj", "dense"]  # Correct modules for Phi-2
+
+# Initialize model and tokenizer
+model = None
+tokenizer = None
+fine_tuned_model = None
+document_text = ""  # Store document content for context
+
+def load_base_model() -> str:
+    """Load Phi-2 with 8-bit quantization instead of 4-bit for faster training"""
+    global model, tokenizer
+
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        gc.collect()
+
+    try:
+        # Use 8-bit quantization (faster to train than 4-bit)
+        if DEVICE == "cuda":
+            bnb_config = BitsAndBytesConfig(
+                load_in_8bit=True,
+                llm_int8_threshold=6.0,
+                llm_int8_has_fp16_weight=False
+            )
+        else:
+            bnb_config = None
+
+        # Load tokenizer with Phi-2 specific settings
+        print("Loading Phi-2 tokenizer...")
+        tokenizer = AutoTokenizer.from_pretrained(
+            MODEL_KEY,
+            trust_remote_code=True,
+            padding_side="right"
+        )
+
+        # Ensure pad token is properly set
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+
+        # Load model with Phi-2 specific configuration
+        print("Loading Phi-2 model... (this may take a few minutes)")
+        if DEVICE == "cuda":
+            model = AutoModelForCausalLM.from_pretrained(
+                MODEL_KEY,
+                quantization_config=bnb_config,
+                device_map="auto",
+                torch_dtype=torch.float16,
+                trust_remote_code=True,
+                low_cpu_mem_usage=True
+            )
+        else:
+            model = AutoModelForCausalLM.from_pretrained(
+                MODEL_KEY,
+                torch_dtype=torch.float32,
+                trust_remote_code=True,
+                low_cpu_mem_usage=True
+            ).to(DEVICE)
+
+        print("Phi-2 (2.7B) model loaded successfully!")
+        return "Phi-2 (2.7B) model loaded successfully! Ready to process documents."
+
+    except Exception as e:
+        error_msg = f"Error loading model: {str(e)}"
+        print(error_msg)
+        return error_msg
+
+def phi2_prompt_template(context: str, question: str) -> str:
+    """
+    Create a prompt optimized for Phi-2
+    Phi-2 responds well to clear instruction formatting
+    """
+    return f"""Instruction: Answer the question accurately based on the context provided.
+Context: {context}
+Question: {question}
+Answer:"""
+
+def process_pdf(file_path: str) -> str:
+    """Extract text from PDF file"""
+    text = ""
+    try:
+        with open(file_path, 'rb') as file:
+            pdf_reader = PyPDF2.PdfReader(file)
+            total_pages = len(pdf_reader.pages)
+            # Process at most 30 pages to avoid memory issues
+            pages_to_process = min(total_pages, 30)
+            for i in range(pages_to_process):
+                page = pdf_reader.pages[i]
+                page_text = page.extract_text() or ""
+                text += page_text + "\n"
+
+            if total_pages > pages_to_process:
+                text += f"\n[Note: Only the first {pages_to_process} pages were processed due to size limitations.]"
+    except Exception as e:
+        print(f"Error processing PDF: {str(e)}")
+    return text
+
+def process_docx(file_path: str) -> str:
+    """Extract text from DOCX file"""
+    try:
+        doc = docx.Document(file_path)
+        text = "\n".join([para.text for para in doc.paragraphs])
+        return text
+    except Exception as e:
+        print(f"Error processing DOCX: {str(e)}")
+        return ""
+
+def process_txt(file_path: str) -> str:
+    """Extract text from TXT file"""
+    try:
+        with open(file_path, 'r', encoding='utf-8', errors='ignore') as file:
+            text = file.read()
+        return text
+    except Exception as e:
+        print(f"Error processing TXT: {str(e)}")
+        return ""
+
+def preprocess_text(text: str) -> str:
+    """Clean and preprocess text"""
+    if not text:
+        return ""
+    # Remove extra whitespace
+    text = re.sub(r'\s+', ' ', text)
+    # Remove special characters that may cause issues
+    text = re.sub(r'[^\w\s.,;:!?\'\"()-]', '', text)
+    return text.strip()
+
+def get_semantic_chunks(text: str, chunk_size: int = 300, overlap: int = 50) -> List[str]:
+    """More efficient semantic chunking"""
+    if not text:
+        return []
+
+    # Simple sentence splitting for speed
+    sentences = re.split(r'(?<=[.!?])\s+', text)
+    chunks = []
+    current_chunk = []
+    current_length = 0
+
+    for sentence in sentences:
+        words = sentence.split()
+        if current_length + len(words) <= chunk_size:
+            current_chunk.append(sentence)
+            current_length += len(words)
+        else:
+            if current_chunk:
+                chunks.append(' '.join(current_chunk))
+            current_chunk = [sentence]
+            current_length = len(words)
+
+    if current_chunk:
+        chunks.append(' '.join(current_chunk))
+
+    # Limit to just 5 chunks for much faster processing
+    if len(chunks) > 5:
+        indices = np.linspace(0, len(chunks)-1, 5, dtype=int)
+        chunks = [chunks[i] for i in indices]
+
+    return chunks
+
+def create_qa_dataset(document_chunks: List[str]) -> List[Dict[str, str]]:
+    """Create comprehensive QA pairs from document chunks for better fine-tuning"""
+    qa_pairs = []
+
+    # Document-level questions
+    full_text = " ".join(document_chunks[:5])  # Use beginning of document for overview
+    qa_pairs.append({
+        "question": "What is this document about?",
+        "context": full_text,
+        "answer": "Based on my analysis, this document discusses..."  # Empty template for model to learn
+    })
+
+    qa_pairs.append({
+        "question": "Summarize the key points of this document.",
+        "context": full_text,
+        "answer": "The key points of this document are..."
+    })
+
+    # Process each chunk for specific QA pairs
+    for i, chunk in enumerate(document_chunks):
+        if not chunk or len(chunk) < 100:  # Skip very short chunks
+            continue
+
+        # Context-specific questions
+        chunk_index = i + 1  # 1-indexed for readability
+
+        # Basic factual questions about chunk content
+        qa_pairs.append({
+            "question": f"What information is contained in section {chunk_index}?",
+            "context": chunk,
+            "answer": f"Section {chunk_index} contains information about..."
+        })
+
+        # Entity-based questions - find names, organizations, technical terms
+        entities = set(re.findall(r'\b[A-Z][a-z]+(?:\s+[A-Z][a-z]+)*\b', chunk))
+        technical_terms = set(re.findall(r'\b[A-Za-z]+-?[A-Za-z]+\b', chunk))
+
+        # Filter to meaningful entities (longer than 3 chars)
+        entities = [e for e in entities if len(e) > 3][:2]  # Limit to 2 entity questions per chunk
+
+        for entity in entities:
+            qa_pairs.append({
+                "question": f"What does the document say about {entity}?",
+                "context": chunk,
+                "answer": f"Regarding {entity}, the document states that..."
+            })
+
+        # Specific content questions
+        sentences = re.split(r'(?<=[.!?])\s+', chunk)
+        key_sentences = [s for s in sentences if len(s.split()) > 8][:2]  # Focus on substantive sentences
+
+        for sentence in key_sentences:
+            # Create question from sentence by identifying subject
+            subject_match = re.search(r'^(The|A|An|This|These|Those|Some|Any|Many|Few|All|Most)?\s*([A-Za-z\s]+?)\s+(is|are|was|were|has|have|had|can|could|will|would|may|might)', sentence, re.IGNORECASE)
+            if subject_match:
+                subject = subject_match.group(2).strip()
+                if len(subject) > 2:
+                    qa_pairs.append({
+                        "question": f"What information is provided about {subject}?",
+                        "context": chunk,
+                        "answer": sentence
+                    })
+
+        # Add relationship questions between concepts
+        if i < len(document_chunks) - 1:
+            next_chunk = document_chunks[i+1]
+            qa_pairs.append({
+                "question": f"How does the information in section {chunk_index} relate to section {chunk_index+1}?",
+                "context": chunk + " " + next_chunk,
+                "answer": f"Section {chunk_index} discusses... while section {chunk_index+1} covers... The relationship between them is..."
+            })
+
+    # Limit to 5 examples max for lighter memory usage
+    if len(qa_pairs) > 5:
+        import random
+        random.shuffle(qa_pairs)
+        qa_pairs = qa_pairs[:5]
+
+    return qa_pairs
+
+class QADataset(TorchDataset):
+    """PyTorch dataset specialized for Phi-2 QA fine-tuning"""
+    def __init__(self, qa_pairs: List[Dict[str, str]], tokenizer, max_length: int = MAX_SEQ_LEN):
+        self.qa_pairs = qa_pairs
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+
+        # Verify dataset structure
+        self.validate_dataset()
+
+    def validate_dataset(self):
+        """Verify that the dataset has proper structure"""
+        if not self.qa_pairs:
+            print("Warning: Empty dataset!")
+            return
+
+        required_keys = ["question", "context", "answer"]
+        for i, item in enumerate(self.qa_pairs[:5]):  # Check first 5 examples
+            missing = [k for k in required_keys if k not in item]
+            if missing:
+                print(f"Warning: Example {i} missing keys: {missing}")
+
+            # Check for empty values
+            empty = [k for k in required_keys if k in item and not item[k]]
+            if empty:
+                print(f"Warning: Example {i} has empty values for: {empty}")
+
+    def __len__(self):
+        return len(self.qa_pairs)
+
+    def __getitem__(self, idx):
+        qa_pair = self.qa_pairs[idx]
+
+        # Format prompt using Phi-2 template
+        context = qa_pair['context']
+        question = qa_pair['question']
+        answer = qa_pair['answer']
+
+        # Build Phi-2 specific prompt
+        prompt = phi2_prompt_template(context, question)
+
+        # Concatenate prompt and answer
+        sequence = f"{prompt} {answer}"
+
+        try:
+            # Tokenize with proper handling
+            encoded = self.tokenizer(
+                sequence,
+                truncation=True,
+                max_length=self.max_length,
+                padding="max_length",
+                return_tensors="pt"
+            )
+
+            # Extract tensors
+            input_ids = encoded["input_ids"].squeeze(0)
+            attention_mask = encoded["attention_mask"].squeeze(0)
+
+            # Create labels
+            labels = input_ids.clone()
+
+            # Calculate prompt length accurately
+            prompt_encoded = self.tokenizer(prompt, add_special_tokens=False)
+            prompt_length = len(prompt_encoded["input_ids"])
+
+            # Ensure prompt_length doesn't exceed labels length
+            prompt_length = min(prompt_length, len(labels))
+
+            # Set labels for prompt portion to -100 (ignored in loss calculation)
+            labels[:prompt_length] = -100
+
+            return {
+                "input_ids": input_ids,
+                "attention_mask": attention_mask,
+                "labels": labels
+            }
+
+        except Exception as e:
+            print(f"Error processing sample {idx}: {e}")
+            # Return dummy sample as fallback
+            return {
+                "input_ids": torch.zeros(self.max_length, dtype=torch.long),
+                "attention_mask": torch.zeros(self.max_length, dtype=torch.long),
+                "labels": torch.zeros(self.max_length, dtype=torch.long)
+            }
+
+def clear_gpu_memory():
+    """Clear GPU memory to prevent OOM errors"""
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        gc.collect()
+
+class ProgressCallback(TrainerCallback):
+    def __init__(self, progress, status_box=None):
+        self.progress = progress
+        self.status_box = status_box
+        self.current_step = 0
+        self.total_steps = 0
+
+    def on_train_begin(self, args, state, control, **kwargs):
+        self.total_steps = state.max_steps
+
+    def on_step_end(self, args, state, control, **kwargs):
+        self.current_step = state.global_step
+        progress_percent = self.current_step / self.total_steps
+        self.progress(0.4 + (0.5 * progress_percent),
+                     desc=f"Epoch {state.epoch}/{args.num_train_epochs} | Step {self.current_step}/{self.total_steps}")
+        if self.status_box:
+            self.status_box.update(f"Training in progress: Epoch {state.epoch}/{args.num_train_epochs} | Step {self.current_step}/{self.total_steps}")
+
+def create_deepspeed_config():
+    """Create DeepSpeed config for faster training"""
+    return {
+        "fp16": {
+            "enabled": True
+        },
+        "zero_optimization": {
+            "stage": 2,
+            "offload_optimizer": {
+                "device": "cpu",
+                "pin_memory": True
+            },
+            "allgather_partitions": True,
+            "allgather_bucket_size": 5e8,
+            "reduce_scatter": True,
+            "reduce_bucket_size": 5e8,
+            "overlap_comm": True,
+            "contiguous_gradients": True
+        },
+        "optimizer": {
+            "type": "AdamW",
+            "params": {
+                "lr": 2e-4,
+                "betas": [0.9, 0.999],
+                "eps": 1e-8,
+                "weight_decay": 0.01
+            }
+        },
+        "scheduler": {
+            "type": "WarmupLR",
+            "params": {
+                "warmup_min_lr": 0,
+                "warmup_max_lr": 2e-4,
+                "warmup_num_steps": 50
+            }
+        },
+        "train_batch_size": 1,
+        "train_micro_batch_size_per_gpu": 1,
+        "gradient_accumulation_steps": 1,
+        "gradient_clipping": 0.5,
+        "steps_per_print": 10
+    }
+
+def finetune_model(qa_dataset, progress=gr.Progress(), status_box=None):
+    """Fine-tune Phi-2 using optimized LoRA parameters"""
+    global model, tokenizer, fine_tuned_model
+
+    if model is None:
+        return "Please load the base model first."
+
+    if len(qa_dataset) == 0:
+        return "No training data created. Please check your document."
+
+    try:
+        progress(0.1, desc="Preparing model for fine-tuning...")
+        if status_box:
+            status_box.update("Preparing model for fine-tuning...")
+
+        # Clear GPU memory
+        clear_gpu_memory()
+
+        # Prepare model for 8-bit training if using GPU
+        if DEVICE == "cuda":
+            training_model = prepare_model_for_kbit_training(model)
+        else:
+            training_model = model
+
+        # Add this line to fix the gradient error
+        training_model.enable_input_require_grads()
+
+        # Configure LoRA for Phi-2
+        peft_config = LoraConfig(
+            r=2,  # Reduced rank for lighter training
+            lora_alpha=4,  # Reduced alpha
+            lora_dropout=0.05,  # Added small dropout for regularization
+            bias="none",
+            task_type="CAUSAL_LM",
+            target_modules=LORA_TARGET_MODULES  # Fixed Phi-2 modules
+        )
+
+        # Apply LoRA to model
+        lora_model = get_peft_model(training_model, peft_config)
+
+        # Print trainable parameters
+        trainable_params = sum(p.numel() for p in lora_model.parameters() if p.requires_grad)
+        all_params = sum(p.numel() for p in lora_model.parameters())
+        print(f"Trainable parameters: {trainable_params:,} ({trainable_params/all_params:.2%} of {all_params:,} total)")
+
+        # Enable gradient checkpointing for memory efficiency
+        if hasattr(lora_model, "gradient_checkpointing_enable"):
+            lora_model.gradient_checkpointing_enable()
+            print("Gradient checkpointing enabled")
+
+        # Create training arguments optimized for Phi-2
+        training_args = TrainingArguments(
+            output_dir="./results",
+            num_train_epochs=2,  # Set to 2 as requested
+            per_device_train_batch_size=1,
+            gradient_accumulation_steps=1,
+            learning_rate=1e-4,  # Reduced from 2e-4 for stability
+            lr_scheduler_type="constant",  # Simplified scheduler
+            warmup_ratio=0.05,  # Slight increase in warmup
+            weight_decay=0.01,
+            logging_steps=1,
+            max_grad_norm=0.3,  # Reduced from 0.5 for better gradient stability
+            save_strategy="no",
+            report_to="none",
+            remove_unused_columns=False,
+            fp16=(DEVICE == "cuda"),
+            no_cuda=(DEVICE == "cpu"),
+            optim="adamw_torch",  # Use standard optimizer instead of fused for stability
+            gradient_checkpointing=True
+        )
+
+        # Add DeepSpeed if on CUDA
+        if DEVICE == "cuda":
+            training_args.deepspeed = create_deepspeed_config()
+
+        # Create data collator that doesn't move tensors to device yet
+        def collate_fn(features):
+            batch = {}
+            for key in features[0].keys():
+                if key in ["input_ids", "attention_mask", "labels"]:
+                    batch[key] = torch.stack([f[key] for f in features])
+            return batch
+
+        progress(0.3, desc="Setting up trainer...")
+        if status_box:
+            status_box.update("Setting up trainer...")
+
+        # Create trainer
+        trainer = Trainer(
+            model=lora_model,
+            args=training_args,
+            train_dataset=qa_dataset,
+            data_collator=collate_fn,
+            callbacks=[ProgressCallback(progress, status_box)]  # Add both callbacks
+        )
+
+        # Start training
+        progress(0.4, desc="Initializing training...")
+        if status_box:
+            status_box.update("Initializing training...")
+        print("Starting training...")
+        trainer.train()
+
+        # Set fine-tuned model
+        fine_tuned_model = lora_model
+
+        # Put model in evaluation mode
+        fine_tuned_model.eval()
+
+        # Clear memory
+        clear_gpu_memory()
+
+        return "Fine-tuning completed successfully! You can now ask questions about your document."
+
+    except Exception as e:
+        error_msg = f"Error during fine-tuning: {str(e)}"
+        print(error_msg)
+        import traceback
+        traceback.print_exc()
+
+        # Try to clean up memory
+        try:
+            clear_gpu_memory()
+        except:
+            pass
+
+        return error_msg
+
+def process_document(file_obj, progress=gr.Progress(), status_box=None):
+    """Process uploaded document and prepare dataset for fine-tuning"""
+    global model, tokenizer, document_text
+
+    progress(0, desc="Processing document...")
+    if status_box:
+        status_box.update("Processing document...")
+
+    if not file_obj:
+        return "Please upload a document first."
+
+    try:
+        # Create temp directory for file
+        temp_dir = tempfile.mkdtemp()
+
+        # Get file name
+        file_name = getattr(file_obj, 'name', 'uploaded_file')
+        if not isinstance(file_name, str):
+            file_name = "uploaded_file.txt"  # Default name
+
+        # Ensure file has extension
+        if '.' not in file_name:
+            file_name = file_name + '.txt'
+
+        temp_path = os.path.join(temp_dir, file_name)
+
+        # Get file content
+        if hasattr(file_obj, 'read'):
+            file_content = file_obj.read()
+        else:
+            file_content = file_obj
+
+        with open(temp_path, 'wb') as f:
+            f.write(file_content)
+
+        # Extract text based on file extension
+        file_extension = os.path.splitext(file_name)[1].lower()
+
+        if file_extension == '.pdf':
+            text = process_pdf(temp_path)
+        elif file_extension in ['.docx', '.doc']:
+            text = process_docx(temp_path)
+        elif file_extension == '.txt' or True:  # Default to txt for unknown extensions
+            text = process_txt(temp_path)
+
+        # Check if text was extracted
+        if not text or len(text) < 50:
+            return "Could not extract sufficient text from the document. Please check the file."
+
+        # Save document text for context window during inference
+        document_text = text
+
+        # Preprocess and chunk the document
+        progress(0.3, desc="Preprocessing document...")
+        if status_box:
+            status_box.update("Preprocessing document...")
+        text = preprocess_text(text)
+        chunks = get_semantic_chunks(text)
+
+        if not chunks:
+            return "Could not extract meaningful text from the document."
+
+        # Create enhanced QA pairs
+        progress(0.5, desc="Creating QA dataset...")
+        if status_box:
+            status_box.update("Creating QA dataset...")
+        qa_pairs = create_qa_dataset(chunks)
+
+        print(f"Created {len(qa_pairs)} QA pairs for training")
+
+        # Debug: Print a sample of QA pairs to verify format
+        if qa_pairs:
+            print("\nSample QA pair for validation:")
+            sample = qa_pairs[0]
+            print(f"Question: {sample['question']}")
+            print(f"Context length: {len(sample['context'])} chars")
+            print(f"Answer: {sample['answer'][:50]}...")
+
+        # Create dataset
+        qa_dataset = QADataset(qa_pairs, tokenizer, max_length=MAX_SEQ_LEN)
+
+        # Fine-tune model
+        progress(0.7, desc="Starting fine-tuning...")
+        if status_box:
+            status_box.update("Starting fine-tuning...")
+        result = finetune_model(qa_dataset, progress, status_box)
+
+        # Clean up
+        try:
+            os.remove(temp_path)
+            os.rmdir(temp_dir)
+        except:
+            pass
+
+        return result
+
+    except Exception as e:
+        error_msg = f"Error processing document: {str(e)}"
+        print(error_msg)
+        import traceback
+        traceback.print_exc()
+        return error_msg
+
+def generate_answer(question, status_box=None):
+    """Generate answer using fine-tuned Phi-2 model with improved response quality"""
+    global fine_tuned_model, tokenizer, document_text
+
+    if fine_tuned_model is None:
+        return "Please process a document first!"
+
+    if not question.strip():
+        return "Please enter a question."
+
+    try:
+        # Clear memory before generation
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+        # For better answers, use document context to help the model
+        # Find relevant context from document (simple keyword matching for efficiency)
+        keywords = re.findall(r'\b\w{5,}\b', question.lower())
+        context = document_text
+
+        # If document is very long, try to find relevant section
+        if len(document_text) > 2000 and keywords:
+            chunks = get_semantic_chunks(document_text, chunk_size=500, overlap=100)
+            relevant_chunks = []
+
+            for chunk in chunks:
+                score = sum(1 for keyword in keywords if keyword.lower() in chunk.lower())
+                if score > 0:
+                    relevant_chunks.append((chunk, score))
+
+            relevant_chunks.sort(key=lambda x: x[1], reverse=True)
+
+            if relevant_chunks:
+                # Use top 2 most relevant chunks
+                context = " ".join([chunk for chunk, _ in relevant_chunks[:2]])
+
+        # Limit context length to fit in model's context window
+        context = context[:1500]  # Limit to 1500 chars for prompt space
+
+        # Create Phi-2 optimized prompt
+        prompt = phi2_prompt_template(context, question)
+
+        # Ensure model is in evaluation mode
+        fine_tuned_model.eval()
+
+        # Tokenize input
+        inputs = tokenizer(prompt, return_tensors="pt").to(fine_tuned_model.device)
+
+        # Configure generation parameters optimized for Phi-2
+        with torch.no_grad():
+            outputs = fine_tuned_model.generate(
+                **inputs,
+                max_new_tokens=75,  # Reduced from 150
+                do_sample=True,
+                temperature=0.7,
+                top_k=40,
+                top_p=0.85,
+                repetition_penalty=1.2,
+                pad_token_id=tokenizer.pad_token_id
+            )
+
+        # Decode response
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+        # Extract only the generated answer part
+        if "Answer:" in response:
+            answer = response.split("Answer:")[-1].strip()
+        else:
+            answer = response
+
+        # If answer is too short or generic, try again with more temperature
+        if len(answer.split()) < 10 or "I don't have enough information" in answer:
+            with torch.no_grad():
+                outputs = fine_tuned_model.generate(
+                    **inputs,
+                    max_new_tokens=75,  # Reduced from 150
+                    do_sample=True,
+                    temperature=0.9,  # Higher temperature
+                    top_k=40,
+                    top_p=0.92,
+                    repetition_penalty=1.2,
+                    pad_token_id=tokenizer.pad_token_id
+                )
+
+            # Decode second attempt
+            response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+            # Extract answer
+            if "Answer:" in response:
+                answer = response.split("Answer:")[-1].strip()
+            else:
+                answer = response
+
+        return answer
+
+    except Exception as e:
+        error_msg = f"Error generating answer: {str(e)}"
+        print(error_msg)
+        return error_msg
+
+# Create Gradio interface
+with gr.Blocks(title="Phi-2 Document QA", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 📚 Phi-2 Document Q&A System")
+    gr.Markdown("Specialized system for fine-tuning Microsoft's Phi-2 model on your documents")
+
+    with gr.Tab("Document Processing"):
+        file_input = gr.File(
+            label="Upload Document (PDF, DOCX, or TXT)",
+            file_types=[".pdf", ".docx", ".txt"],
+            type="binary"
+        )
+
+        with gr.Row():
+            load_model_btn = gr.Button("1. Load Phi-2 Model", variant="secondary")
+            process_btn = gr.Button("2. Process & Fine-tune Document", variant="primary")
+
+        status = gr.Textbox(
+            label="Status",
+            placeholder="First load the model, then upload a document and click 'Process & Fine-tune'",
+            lines=3
+        )
+
+        gr.Markdown("""
+        ### Tips for Best Results
+        - PDF, DOCX and TXT files are supported
+        - Keep documents under 10 pages for best results
+        - Processing time depends on document length and GPU availability
+        - For GPU usage in Colab: Runtime > Change runtime type > GPU
+        """)
+
+    with gr.Tab("Ask Questions"):
+        question_input = gr.Textbox(
+            label="Your Question",
+            placeholder="Ask about your document...",
+            lines=2
+        )
+
+        ask_btn = gr.Button("Get Answer", variant="primary")
+
+        answer_output = gr.Textbox(
+            label="Phi-2's Response",
+            placeholder="The answer will appear here after you ask a question",
+            lines=8
+        )
+
+        gr.Markdown("""
+        ### Example Questions
+        - "What is this document about?"
+        - "Summarize the key points in this document"
+        - "What does the document say about [specific topic]?"
+        - "Explain the relationship between [concept A] and [concept B]"
+        """)
+
+    # Set up events
+    load_model_btn.click(
+        fn=load_base_model,
+        outputs=[status]
+    )
+
+    process_btn.click(
+        fn=process_document,
+        inputs=[file_input],
+        outputs=[status]
+    )
+
+    ask_btn.click(
+        fn=generate_answer,
+        inputs=[question_input],
+        outputs=[answer_output]
+    )
+
+# Launch the app
+if __name__ == "__main__":
+    demo.launch(share=True)