Update app.py

app.py

@@ -1,5 +1,11 @@
+import transformers
+import pickle
 import os
 import numpy as np
+import torchvision
+import nltk
+import torch
+import pandas as pd
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
@@ -8,17 +14,17 @@ from transformers import (
     AutoModelForSeq2SeqLM,
     AutoModelForTokenClassification,
     AutoModelForCausalLM,
-    pipeline
+    pipeline, 
+    Qwen2Tokenizer, 
+    BartForConditionalGeneration
 )
-from sentence_transformers import SentenceTransformer, CrossEncoder
+from sentence_transformers import SentenceTransformer, CrossEncoder, util
 from sklearn.metrics.pairwise import cosine_similarity
 from bs4 import BeautifulSoup
-import nltk
-import torch
-import pandas as pd
 from huggingface_hub import hf_hub_download
-from safetensors.torch import load_file  # Import Safetensors loader
+from safetensors.torch import load_file 
 from typing import List, Dict, Optional
+from safetensors.numpy import load_file
 
 # Initialize FastAPI app
 app = FastAPI()
@@ -77,6 +83,7 @@ def load_models():
         # Embedding models
         models['embedding'] = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
         models['cross_encoder'] = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2', max_length=512)
+        models['semantic_model'] = SentenceTransformer('all-MiniLM-L6-v2')
 
         # Translation models
         models['ar_to_en_tokenizer'] = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-ar-en")
@@ -84,6 +91,10 @@ def load_models():
         models['en_to_ar_tokenizer'] = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-ar")
         models['en_to_ar_model'] = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-en-ar")
 
+        #Attention model
+        models['att_tokenizer'] = AutoTokenizer.from_pretrained("facebook/bart-base")
+        models['att_model'] = BartForConditionalGeneration.from_pretrained("facebook/bart-base")
+        
         # NER model
         models['bio_tokenizer'] = AutoTokenizer.from_pretrained("blaze999/Medical-NER")
         models['bio_model'] = AutoModelForTokenClassification.from_pretrained("blaze999/Medical-NER")
@@ -100,41 +111,71 @@ def load_models():
         print(f"Error loading models: {e}")
         return False
 
-
 def load_embeddings() -> Optional[Dict[str, np.ndarray]]:
     """Load embeddings from Safetensors file"""
     try:
+        # Locate or download embeddings file
         embeddings_path = 'embeddings.safetensors'
         if not os.path.exists(embeddings_path):
+            print("File not found locally. Attempting to download from Hugging Face Hub...")
             embeddings_path = hf_hub_download(
-                repo_id=os.environ.get('thechaiexperiment/TeaRAG', ''),
+                repo_id=os.environ.get('HF_SPACE_ID', 'thechaiexperiment/TeaRAG'),
                 filename="embeddings.safetensors",
                 repo_type="space"
             )
-
+        
+        # Load Safetensors file
         embeddings = load_file(embeddings_path)
         if not isinstance(embeddings, dict):
-            raise ValueError("Invalid format for embeddings in Safetensors file.")
+            raise ValueError("Expected a dictionary in the Safetensors file.")
+        
+        # Validate and convert tensors to numpy arrays
+        result = {}
+        for key, tensor in embeddings.items():
+            if not hasattr(tensor, 'numpy'):
+                raise TypeError(f"Value for key {key} is not a tensor or cannot be converted to numpy.")
+            result[key] = tensor.numpy()
+        
+        print("Embeddings successfully loaded.")
+        return result
 
-        # Convert to dictionary with numpy arrays
-        return {k: tensor.numpy() for k, tensor in embeddings.items()}
     except Exception as e:
         print(f"Error loading embeddings: {e}")
         return None
 
-        
-def load_documents_data():
-    """Load document data with error handling"""
+def load_documents_data(folder_path='downloaded_articles/downloaded_articles'):
+    """Load document data from HTML articles in a specified folder."""
     try:
         print("Loading documents data...")
-        docs_path = 'finalcleaned_excel_file.xlsx'
+        # Check if the folder exists
+        if not os.path.exists(folder_path) or not os.path.isdir(folder_path):
+            print(f"Error: Folder '{folder_path}' not found")
+            return False
+        # List all HTML files in the folder
+        html_files = [f for f in os.listdir(folder_path) if f.endswith('.html')]
+        if not html_files:
+            print(f"No HTML files found in folder '{folder_path}'")
+            return False
+        documents = []
+        # Iterate through each HTML file and parse the content
+        for file_name in html_files:
+            file_path = os.path.join(folder_path, file_name)
+            try:
+                with open(file_path, 'r', encoding='utf-8') as file:
+                    # Parse the HTML file
+                    soup = BeautifulSoup(file, 'html.parser')
+                    # Extract text content (or customize this as per your needs)
+                    text = soup.get_text(separator='\n').strip()
+                    documents.append({"file_name": file_name, "content": text})
+            except Exception as e:
+                print(f"Error reading file {file_name}: {e}")
+        # Convert the list of documents to a DataFrame
+        data['df'] = pd.DataFrame(documents)
         
-        if not os.path.exists(docs_path):
-            print(f"Error: {docs_path} not found")
+        if data['df'].empty:
+            print("No valid documents loaded.")
             return False
-            
-        data['df'] = pd.read_excel(docs_path)
-        print(f"Successfully loaded {len(data['df'])} document records")
+        print(f"Successfully loaded {len(data['df'])} document records.")
         return True
     except Exception as e:
         print(f"Error loading documents data: {e}")
@@ -174,14 +215,17 @@ def embed_query_text(query_text):
     query_embedding = embedding.encode([query_text])
     return query_embedding
     
-def query_embeddings(query_embedding, n_results=5):
-    """Find relevant documents using embedding similarity"""
-    if not data['embeddings']:
+from sklearn.metrics.pairwise import cosine_similarity
+import numpy as np
+
+def query_embeddings(query_embedding, embeddings_data=None, n_results=5):
+    embeddings_data = embeddings_data or data.get('embeddings', {})
+    if not embeddings_data:
+        print("No embeddings data available.")
         return []
-        
     try:
-        doc_ids = list(data['embeddings'].keys())
-        doc_embeddings = np.array(list(data['embeddings'].values()))
+        doc_ids = list(embeddings_data.keys())
+        doc_embeddings = np.array(list(embeddings_data.values()))
         similarities = cosine_similarity(query_embedding, doc_embeddings).flatten()
         top_indices = similarities.argsort()[-n_results:][::-1]
         return [(doc_ids[i], similarities[i]) for i in top_indices]
@@ -189,66 +233,85 @@ def query_embeddings(query_embedding, n_results=5):
         print(f"Error in query_embeddings: {e}")
         return []
 
-def retrieve_document_text(doc_id):
-    """Retrieve document text from HTML file"""
-    try:
-        file_path = os.path.join('downloaded_articles', doc_id)
-        if not os.path.exists(file_path):
-            print(f"Warning: Document file not found: {file_path}")
-            return ""
-            
-        with open(file_path, 'r', encoding='utf-8') as file:
-            soup = BeautifulSoup(file, 'html.parser')
-            return soup.get_text(separator=' ', strip=True)
-    except Exception as e:
-        print(f"Error retrieving document {doc_id}: {e}")
-        return ""
-
-
-def rerank_documents(query, doc_texts):
-    """Rerank documents using cross-encoder"""
+def retrieve_document_texts(doc_ids, folder_path='downloaded_articles/downloaded_articles'):
+    texts = []
+    for doc_id in doc_ids:
+        file_path = os.path.join(folder_path, doc_id)
+        try:
+            # Check if the file exists
+            if not os.path.exists(file_path):
+                print(f"Warning: Document file not found: {file_path}")
+                texts.append("")
+                continue
+            # Read and parse the HTML file
+            with open(file_path, 'r', encoding='utf-8') as file:
+                soup = BeautifulSoup(file, 'html.parser')
+                text = soup.get_text(separator=' ', strip=True)
+                texts.append(text)
+        except Exception as e:
+            print(f"Error retrieving document {doc_id}: {e}")
+            texts.append("")
+    return texts
+
+
+def rerank_documents(query, document_ids, document_texts, cross_encoder_model):
     try:
-        pairs = [(query, doc) for doc in doc_texts]
-        scores = models['cross_encoder'].predict(pairs)
-        return scores
+        # Prepare pairs for the cross-encoder
+        pairs = [(query, doc) for doc in document_texts]
+        # Get scores using the cross-encoder model
+        scores = cross_encoder_model.predict(pairs)
+        # Combine scores with document IDs and texts
+        scored_documents = list(zip(scores, document_ids, document_texts))
+        # Sort by scores in descending order
+        scored_documents.sort(key=lambda x: x[0], reverse=True)
+        # Print reranked results
+        print("Reranked results:")
+        for idx, (score, doc_id, doc) in enumerate(scored_documents):
+            print(f"Rank {idx + 1} (Score: {score:.4f}, Document ID: {doc_id})")
+        return scored_documents
     except Exception as e:
         print(f"Error reranking documents: {e}")
-        return np.zeros(len(doc_texts))
+        return []
 
-def extract_entities(text):
-    """Extract medical entities from text using NER"""
+def extract_entities(text, ner_pipeline=None):
     try:
-        results = models['ner_pipeline'](text)
-        return list({result['word'] for result in results if result['entity'].startswith("B-")})
+        # Use the provided pipeline or default to the model dictionary
+        if ner_pipeline is None:
+            ner_pipeline = models['ner_pipeline']
+        # Perform NER using the pipeline
+        ner_results = ner_pipeline(text)
+        # Extract unique entities that start with "B-"
+        entities = {result['word'] for result in ner_results if result['entity'].startswith("B-")}
+        return list(entities)
     except Exception as e:
         print(f"Error extracting entities: {e}")
         return []
+
 def match_entities(query_entities, sentence_entities):
-    query_set, sentence_set = set(query_entities), set(sentence_entities)
-    matches = query_set.intersection(sentence_set)
-    return len(matches)
+    try:
+        query_set, sentence_set = set(query_entities), set(sentence_entities)
+        matches = query_set.intersection(sentence_set)
+        return len(matches)
+    except Exception as e:
+        print(f"Error matching entities: {e}")
+        return 0
 
 def extract_relevant_portions(document_texts, query, max_portions=3, portion_size=1, min_query_words=1):
     relevant_portions = {}
-
     # Extract entities from the query
     query_entities = extract_entities(query, ner_biobert)
     print(f"Extracted Query Entities: {query_entities}")
     for doc_id, doc_text in enumerate(document_texts):
         sentences = nltk.sent_tokenize(doc_text)  # Split document into sentences
         doc_relevant_portions = []
-
         # Extract entities from the entire document
         doc_entities = extract_entities(doc_text, ner_biobert)
         print(f"Document {doc_id} Entities: {doc_entities}")
-
         for i, sentence in enumerate(sentences):
             # Extract entities from the sentence
             sentence_entities = extract_entities(sentence, ner_biobert)
-
             # Compute relevance score
             relevance_score = match_entities(query_entities, sentence_entities)
-
             # Select sentences with at least `min_query_words` matching entities
             if relevance_score >= min_query_words:
                 start_idx = max(0, i - portion_size // 2)
@@ -257,40 +320,43 @@ def extract_relevant_portions(document_texts, query, max_portions=3, portion_siz
                 doc_relevant_portions.append(portion)
             if len(doc_relevant_portions) >= max_portions:
                 break
-
-        # Add fallback to include the most entity-dense sentences if no results
+        # Fallback: Include most entity-dense sentences if no relevant portions were found
         if not doc_relevant_portions and len(doc_entities) > 0:
             print(f"Fallback: Selecting sentences with most entities for Document {doc_id}")
             sorted_sentences = sorted(sentences, key=lambda s: len(extract_entities(s, ner_biobert)), reverse=True)
             for fallback_sentence in sorted_sentences[:max_portions]:
                 doc_relevant_portions.append(fallback_sentence)
-
+        # Add the extracted portions to the result dictionary
         relevant_portions[f"Document_{doc_id}"] = doc_relevant_portions
-
     return relevant_portions
-    
 
 def remove_duplicates(selected_parts):
     unique_sentences = set()
     unique_selected_parts = []
-
     for sentence in selected_parts:
         if sentence not in unique_sentences:
             unique_selected_parts.append(sentence)
             unique_sentences.add(sentence)
-
     return unique_selected_parts
 
 def extract_entities(text):
-    inputs = biobert_tokenizer(text, return_tensors="pt")
-    outputs = biobert_model(**inputs)
-    predictions = torch.argmax(outputs.logits, dim=2)
-    tokens = biobert_tokenizer.convert_ids_to_tokens(inputs.input_ids[0])
-    entities = [tokens[i] for i in range(len(tokens)) if predictions[0][i].item() != 0]  # Assume 0 is the label for non-entity
-    return entities
+    try:
+        inputs = biobert_tokenizer(text, return_tensors="pt")
+        outputs = biobert_model(**inputs)
+        predictions = torch.argmax(outputs.logits, dim=2)
+
+        tokens = biobert_tokenizer.convert_ids_to_tokens(inputs.input_ids[0])
+        entities = [
+            tokens[i] 
+            for i in range(len(tokens)) 
+            if predictions[0][i].item() != 0  # Assuming 0 is the label for non-entity
+        ]
+        return entities
+    except Exception as e:
+        print(f"Error extracting entities: {e}")
+        return []
 
 def enhance_passage_with_entities(passage, entities):
-    # Example: Add entities to the passage for better context
     return f"{passage}\n\nEntities: {', '.join(entities)}"
 
 def create_prompt(question, passage):
@@ -304,13 +370,12 @@ def create_prompt(question, passage):
     Answer:
     """)
     return prompt.format(passage=passage, question=question)
-    
+
 def generate_answer(prompt, max_length=860, temperature=0.2):
     inputs = tokenizer_f(prompt, return_tensors="pt", truncation=True)
-    
     # Start timing
     start_time = time.time()
-    
+    # Generate the output
     output_ids = model_f.generate(
         inputs.input_ids,
         max_length=max_length,
@@ -318,28 +383,25 @@ def generate_answer(prompt, max_length=860, temperature=0.2):
         temperature=temperature,
         pad_token_id=tokenizer_f.eos_token_id
     )
-    
     # End timing
     end_time = time.time()
-    
     # Calculate the duration
     duration = end_time - start_time
-    
     # Decode the answer
     answer = tokenizer_f.decode(output_ids[0], skip_special_tokens=True)
-    
-    passage_keywords = set(passage.lower().split())
+    # Extract keywords from the passage and answer
+    passage_keywords = set(prompt.lower().split())  # Adjusted to check keywords in the full prompt
     answer_keywords = set(answer.lower().split())
-    
+    # Verify if the answer aligns with the passage
     if passage_keywords.intersection(answer_keywords):
         return answer, duration
     else:
         return "Sorry, I can't help with that.", duration
-        
+  
 def remove_answer_prefix(text):
     prefix = "Answer:"
     if prefix in text:
-        return text.split(prefix)[-1].strip()
+        return text.split(prefix, 1)[-1].strip()  # Split only once to avoid splitting at other occurrences of "Answer:"
     return text
 
 def remove_incomplete_sentence(text):
@@ -356,7 +418,6 @@ def remove_incomplete_sentence(text):
 async def root():
     return {"message": "Welcome to the FastAPI application! Use the /health endpoint to check health, and /api/query for processing queries."}
 
-
 @app.get("/health")
 async def health_check():
     """Health check endpoint"""