Update app.py

app.py

@@ -33,9 +33,389 @@ if not HF_TOKEN:
     print("Warning: HF_TOKEN not found in environment variables. You may not be able to access gated models.")
 
 class MedicalReportAnalyzer:
-    ...
-    # Keep your existing MedicalReportAnalyzer implementation here
-    # No changes needed
+    def __init__(self):
+        self.vector_store = None
+        self.llm = None
+        self.qa_chain = None
+        self.user_report_data = "No report data available."  # Default value
+        self.original_report_data = "No original report data available."  # Store original data
+        # Initialize everything
+        self._load_or_create_vector_store()
+        self._initialize_llm()
+        self._setup_qa_chain()
+
+    def _load_or_create_vector_store(self):
+        """Load existing vector store or create a new one from knowledge documents"""
+        embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL)
+        
+        # Check if vector store exists
+        if os.path.exists(VECTOR_STORE_PATH):
+            print("Loading existing vector store...")
+            self.vector_store = FAISS.load_local(VECTOR_STORE_PATH, embeddings)
+        else:
+            print("Creating new vector store from documents...")
+            # Create knowledge directory if it doesn't exist
+            os.makedirs(KNOWLEDGE_DIR, exist_ok=True)
+            
+            # Check if there are documents to process
+            if len(os.listdir(KNOWLEDGE_DIR)) == 0:
+                print(f"Warning: No documents found in {KNOWLEDGE_DIR}. Please add medical PDFs.")
+                # Initialize empty vector store
+                self.vector_store = FAISS.from_texts(["No medical knowledge available yet."], embeddings)
+                self.vector_store.save_local(VECTOR_STORE_PATH)
+                return
+            
+            # Load all PDFs from the knowledge directory
+            try:
+                # First try with DirectoryLoader
+                loader = DirectoryLoader(KNOWLEDGE_DIR, glob="**/*.pdf", loader_cls=PyPDFLoader)
+                documents = loader.load()
+                
+                # Split documents into chunks
+                text_splitter = RecursiveCharacterTextSplitter(
+                    chunk_size=1000,
+                    chunk_overlap=200,
+                    length_function=len
+                )
+                chunks = text_splitter.split_documents(documents)
+                
+                # Create and save the vector store
+                self.vector_store = FAISS.from_documents(chunks, embeddings)
+                self.vector_store.save_local(VECTOR_STORE_PATH)
+            except Exception as e:
+                print(f"Error loading documents with DirectoryLoader: {str(e)}")
+                # Initialize with minimal data
+                self.vector_store = FAISS.from_texts(["No medical knowledge available yet."], embeddings)
+                self.vector_store.save_local(VECTOR_STORE_PATH)
+
+    def _initialize_llm(self):
+        """Initialize the language model with HF token authentication"""
+        print(f"Loading model {MODEL_NAME} on {DEVICE}...")
+        try:
+            # Use the HF_TOKEN for authentication
+            tokenizer = AutoTokenizer.from_pretrained(
+                MODEL_NAME,
+                token=HF_TOKEN
+            )
+            model = AutoModelForCausalLM.from_pretrained(
+                MODEL_NAME,
+                token=HF_TOKEN,
+                torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
+                device_map="auto",
+                load_in_8bit=DEVICE == "cuda",  # Use 8-bit quantization if on CUDA
+            )
+            
+            # Create a text generation pipeline
+            pipe = pipeline(
+                "text-generation",
+                model=model,
+                tokenizer=tokenizer,
+                max_new_tokens=512,
+                temperature=0.1,
+                top_p=0.95,
+                repetition_penalty=1.1
+            )
+            
+            # Create LangChain wrapper around the pipeline
+            self.llm = HuggingFacePipeline(pipeline=pipe)
+        except Exception as e:
+            print(f"Error loading the model: {str(e)}")
+            print("Falling back to a non-gated model...")
+            # Fallback to a non-gated model
+            fallback_model = "google/flan-t5-large"
+            tokenizer = AutoTokenizer.from_pretrained(fallback_model)
+            model = AutoModelForCausalLM.from_pretrained(
+                fallback_model,
+                torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
+                device_map="auto"
+            )
+            pipe = pipeline(
+                "text-generation",
+                model=model,
+                tokenizer=tokenizer,
+                max_new_tokens=512
+            )
+            self.llm = HuggingFacePipeline(pipeline=pipe)
+
+    def _setup_qa_chain(self):
+        """Set up the question-answering chain"""
+        # Define a custom prompt template for medical analysis
+        template = """
+        You are a medical assistant analyzing patient medical reports. Use the following pieces of context to answer the question. If you don't know the answer, just say that you don't know, don't try to make up an answer.
+
+        Patient Report Summary: {patient_data}
+
+        Context from medical knowledge base: {context}
+
+        Question: {question}
+
+        Answer:
+        """
+        
+        prompt = PromptTemplate(
+            template=template,
+            input_variables=["context", "question", "patient_data"]
+        )
+        
+        # Create the QA chain
+        self.qa_chain = RetrievalQA.from_chain_type(
+            llm=self.llm,
+            chain_type="stuff",
+            retriever=self.vector_store.as_retriever(search_kwargs={"k": 5}),
+            chain_type_kwargs={"prompt": prompt},
+            return_source_documents=False
+        )
+    
+    def remove_header_information(self, text):
+        """Remove header information from the report text"""
+        # Store the original text
+        self.original_report_data = text
+        
+        # Split the text into lines to analyze
+        lines = text.split('\n')
+        
+        # Define patterns to identify header information
+        header_patterns = [
+            r'(Name\s*:)',
+            r'(Patient\s*Name\s*:)',
+            r'(DOB|Date of Birth\s*:)',
+            r'(Age\s*:)',
+            r'(Gender\s*:)',
+            r'(Lab No\.|Laboratory Number\s*:)',
+            r'(Patient ID\s*:)',
+            r'(Report Status\s*:)',
+            r'(Ref By|Referred By\s*:)',
+            r'(Collected\s*:)',
+            r'(Reported\s*:)',
+            r'(A/c Status\s*:)',
+            r'(Processed at\s*:)',
+            r'(Collected at\s*:)',
+            r'(Address\s*:)',
+            r'(Phone|Mobile|Mob\s*:)',
+        ]
+        
+        # Create a regex pattern that matches any of the header patterns
+        combined_pattern = '|'.join(header_patterns)
+        
+        # Find where the actual test results begin
+        test_results_start = -1
+        for i, line in enumerate(lines):
+            if re.search(r'(Test\s*Report|Test\s*Name|Test\s*Results|Results|HEMOGRAM|ROUTINE|EXAMINATION)', line, re.IGNORECASE):
+                test_results_start = i
+                break
+        
+        # If we couldn't find the start of test results, look for key medical terms
+        if test_results_start == -1:
+            for i, line in enumerate(lines):
+                # Look for common test result sections
+                if re.search(r'(Hemoglobin|Blood|Urine|CBC|Glucose|Cholesterol|Protein|RBC|WBC)', line, re.IGNORECASE):
+                    test_results_start = max(0, i-3)  # Start a few lines before the first test result
+                    break
+        
+        # If we still couldn't find the start of test results, use a heuristic:
+        # Skip the first few lines which usually contain header information
+        if test_results_start == -1:
+            # Count lines with patient identifiable information
+            header_count = 0
+            for i, line in enumerate(lines):
+                if re.search(combined_pattern, line, re.IGNORECASE):
+                    header_count += 1
+                    
+            # If we found several header lines, skip those plus a few more
+            if header_count > 0:
+                test_results_start = min(header_count + 5, len(lines) // 3)
+            else:
+                # If no clear header pattern was found, just skip the first 10% of lines as a fallback
+                test_results_start = max(1, len(lines) // 10)
+                
+        # Return text from the determined start point
+        clean_text = '\n'.join(lines[test_results_start:])
+        
+        # If this dramatically shortened the text, use a less aggressive approach
+        if len(clean_text) < len(text) * 0.5:
+            print("Warning: Header removal may have removed too much content. Using alternative approach.")
+            # Alternative approach: Just remove lines with header patterns
+            filtered_lines = []
+            for line in lines:
+                if not re.search(combined_pattern, line, re.IGNORECASE):
+                    filtered_lines.append(line)
+            clean_text = '\n'.join(filtered_lines)
+            
+        return clean_text
+
+    def extract_text_from_pdf_pymupdf(self, pdf_path):
+        """Extract text from PDF using PyMuPDF (more robust than PyPDF)"""
+        text = ""
+        try:
+            doc = fitz.open(pdf_path)
+            for page in doc:
+                text += page.get_text()
+            doc.close()
+            return text
+        except Exception as e:
+            print(f"PyMuPDF extraction error: {str(e)}")
+            return None
+
+    def extract_text_from_pdf_pypdf(self, pdf_path):
+        """Extract text using PyPDF as a backup method"""
+        try:
+            loader = PyPDFLoader(pdf_path)
+            pages = loader.load()
+            return "\n".join([page.page_content for page in pages])
+        except Exception as e:
+            print(f"PyPDF extraction error: {str(e)}")
+            return None
+
+    def process_user_report(self, report_file):
+        """Process the uploaded medical report with multiple fallback methods"""
+        if report_file is None:
+            return "No file uploaded. Please upload a medical report."
+        
+        # Ensure the uploaded file is read as bytes
+        temp_dir = tempfile.mkdtemp()
+        try:
+            # Copy the uploaded file to the temp directory
+            temp_file_path = os.path.join(temp_dir, "user_report.pdf")
+            
+            # Handle file based on its type
+            try:
+                if isinstance(report_file, str):  # If it's a file path
+                    shutil.copy(report_file, temp_file_path)
+                elif hasattr(report_file, 'name'):  # Gradio file object
+                    with open(temp_file_path, 'wb') as f:
+                        with open(report_file.name, 'rb') as source:
+                            f.write(source.read())
+                else:  # Try to handle as bytes or file-like object
+                    with open(temp_file_path, 'wb') as f:
+                        f.write(report_file.read() if hasattr(report_file, 'read') else report_file)
+            except Exception as e:
+                print(f"Error saving file: {str(e)}")
+                return f"Error saving the uploaded file: {str(e)}"
+                
+            # Try multiple methods to extract text from the PDF
+            text = None
+            
+            # Method 1: PyMuPDF
+            text = self.extract_text_from_pdf_pymupdf(temp_file_path)
+            
+            # Method 2: PyPDF as fallback
+            if not text:
+                text = self.extract_text_from_pdf_pypdf(temp_file_path)
+                
+            # Method 3: Last resort - try to read as raw text
+            if not text:
+                try:
+                    with open(temp_file_path, 'r', errors='ignore') as f:
+                        text = f.read()
+                except Exception as e:
+                    print(f"Raw text reading error: {str(e)}")
+            
+            # If we got text, process it
+            if text and len(text.strip()) > 0:
+                # Remove header information from the text
+                cleaned_text = self.remove_header_information(text)
+                
+                # Store the cleaned text
+                self.user_report_data = cleaned_text
+                
+                # Split into chunks if needed
+                text_splitter = RecursiveCharacterTextSplitter(
+                    chunk_size=1000,
+                    chunk_overlap=200,
+                    length_function=len
+                )
+                chunks = text_splitter.split_text(cleaned_text)
+                
+                # Check if too much text was removed
+                original_length = len(text.strip())
+                cleaned_length = len(cleaned_text.strip())
+                removal_percentage = (original_length - cleaned_length) / original_length * 100
+                
+                if removal_percentage > 80:
+                    return f"Report processed successfully, but significant content may have been filtered. Original length: {original_length} chars. Cleaned length: {cleaned_length} chars. Extracted approximately {len(chunks)} text chunks."
+                else:
+                    return f"Report processed successfully. Removed approximately {removal_percentage:.1f}% of header content. Extracted {len(chunks)} text chunks."
+            else:
+                self.user_report_data = "Unable to extract text from the provided PDF. This is an empty report placeholder."
+                return "Warning: Could not extract text from the PDF. The file may be corrupted, password-protected, or contain only images. Processing will continue with limited data."
+                
+        finally:
+            # Clean up the temporary directory and file
+            shutil.rmtree(temp_dir)
+
+    def answer_question(self, question):
+        """Answer a question based on the uploaded report and knowledge base"""
+        if not self.user_report_data or self.user_report_data == "No report data available.":
+            return "No report has been processed or text extraction failed. Please upload a medical report first."
+        
+        # Get context from knowledge base
+        try:
+            retrieved_docs = self.vector_store.similarity_search(question, k=5)
+            context = "\n\n".join([doc.page_content for doc in retrieved_docs])
+            
+            # Check if question is about patient demographics or identification
+            demographic_patterns = [
+                r'(patient|name|age|gender|birth|dob|address|phone|contact|id|identification)',
+                r'(doctor|physician|referring|referred by)',
+                r'(date|time|collected|processed|reported)',
+                r'(lab|laboratory|number|id)'
+            ]
+            
+            combined_demo_pattern = '|'.join(demographic_patterns)
+            
+            # If question might be about demographics, check if we need to use original data
+            if re.search(combined_demo_pattern, question, re.IGNORECASE):
+                # For demographic questions, we can use the original report that includes headers
+                # But only if we have specific identification information requests
+                specific_id_patterns = [
+                    r'(name of|patient name|who is|what is the name)',
+                    r'(exact age|age of|how old)',
+                    r'(address of|where|location|contact details)',
+                    r'(doctor name|name of doctor|referring doctor|who referred)',
+                    r'(date of|when was|time of|report date)',
+                    r'(lab number|patient id|identification number)'
+                ]
+                
+                specific_id_pattern = '|'.join(specific_id_patterns)
+                
+                # If it's a direct question about patient identity, don't answer
+                if re.search(specific_id_pattern, question, re.IGNORECASE):
+                    return "I'm unable to provide specific patient identification information. This feature is disabled to protect patient privacy. Please ask about medical test results or interpretations instead."
+            
+            # Create the inputs dict for the QA chain
+            inputs = {
+                "query": question,
+                "context": context,
+                "patient_data": self.user_report_data
+            }
+            
+            # Run the chain with the correct parameter structure
+            result = self.qa_chain(inputs)
+            
+            # Extract the answer from the result
+            if isinstance(result, dict) and 'result' in result:
+                return result['result']
+            else:
+                return str(result)
+                
+        except Exception as e:
+            print(f"Error answering question: {str(e)}")
+            error_msg = f"Error processing your question: {str(e)}."
+            
+            # Try direct LLM call as fallback
+            try:
+                direct_prompt = f"""
+                Question about medical report: {question}
+                
+                Patient data available: {self.user_report_data[:800]}... (truncated)
+                
+                Please answer based on this information:
+                """
+                
+                direct_result = self.llm(direct_prompt)
+                return f"{error_msg} Fallback answer: {direct_result}"
+            except:
+                return f"{error_msg} Please try a different question or report."
 
 # Initialize the analyzer
 analyzer = MedicalReportAnalyzer()