Update app.py

app.py

@@ -6,8 +6,9 @@ from flask_cors import CORS
 from transformers import (
     AutoTokenizer, 
     AutoModelForSeq2SeqLM, 
+    AutoModelForTokenClassification,
     AutoModelForCausalLM,
-    AutoModelForTokenClassification
+    pipeline
 )
 from sentence_transformers import SentenceTransformer, CrossEncoder
 from sklearn.metrics.pairwise import cosine_similarity
@@ -15,71 +16,65 @@ from bs4 import BeautifulSoup
 import nltk
 import torch
 import pandas as pd
-from startup import setup_files
-
 
 app = Flask(__name__)
 CORS(app)
-# Environment variables for file paths
-EMBEDDINGS_PATH = os.environ.get('EMBEDDINGS_PATH', 'data/embeddings.pkl')
-LINKS_PATH = os.environ.get('LINKS_PATH', 'data/finalcleaned_excel_file.xlsx')
 
-def init_app():
-    # Download and extract files if they don't exist
-    if not os.path.exists('downloaded_articles'):
-        setup_files()
+# Global variables for models and data
+models = {}
+data = {}
 
-# Initialize models with proper error handling
-def initialize_models():
+def init_nltk():
+    """Initialize NLTK resources"""
     try:
-        global embedding_model, cross_encoder, semantic_model
-        global ar_to_en_tokenizer, ar_to_en_model
-        global en_to_ar_tokenizer, en_to_ar_model
-        global tokenizer_f, model_f, bio_tokenizer, bio_model
+        nltk.download('punkt', quiet=True)
+        return True
+    except Exception as e:
+        print(f"Error initializing NLTK: {e}")
+        return False
 
-        print("Initializing models...")
+def load_models():
+    """Initialize all required models"""
+    try:
+        print("Loading 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)
         
-        # Basic embedding models
-        embedding_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
-        cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2', max_length=512)
-        semantic_model = SentenceTransformer('all-MiniLM-L6-v2')
-
         # Translation models
-        ar_to_en_tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-ar-en")
-        ar_to_en_model = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-ar-en")
-        en_to_ar_tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-ar")
-        en_to_ar_model = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-en-ar")
-
-        # Medical NER model
-        bio_tokenizer = AutoTokenizer.from_pretrained("blaze999/Medical-NER")
-        bio_model = AutoModelForTokenClassification.from_pretrained("blaze999/Medical-NER")
-
+        models['ar_to_en_tokenizer'] = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-ar-en")
+        models['ar_to_en_model'] = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-ar-en")
+        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")
+        
+        # NER model
+        models['bio_tokenizer'] = AutoTokenizer.from_pretrained("blaze999/Medical-NER")
+        models['bio_model'] = AutoModelForTokenClassification.from_pretrained("blaze999/Medical-NER")
+        models['ner_pipeline'] = pipeline("ner", model=models['bio_model'], tokenizer=models['bio_tokenizer'])
+        
         # LLM model
         model_name = "M4-ai/Orca-2.0-Tau-1.8B"
-        tokenizer_f = AutoTokenizer.from_pretrained(model_name)
-        model_f = AutoModelForCausalLM.from_pretrained(model_name)
-
-        nltk.download('punkt', quiet=True)
+        models['llm_tokenizer'] = AutoTokenizer.from_pretrained(model_name)
+        models['llm_model'] = AutoModelForCausalLM.from_pretrained(model_name)
         
-        print("Models initialized successfully")
+        print("Models loaded successfully")
         return True
     except Exception as e:
-        print(f"Error initializing models: {e}")
+        print(f"Error loading models: {e}")
         return False
 
-# Load data with error handling
 def load_data():
+    """Load embeddings and document data"""
     try:
-        global embeddings_data, df
-        
-        print("Loading data files...")
+        print("Loading data...")
         
         # Load embeddings
-        with open(EMBEDDINGS_PATH, 'rb') as file:
-            embeddings_data = pickle.load(file)
+        with open('embeddings.pkl', 'rb') as f:
+            data['embeddings'] = pickle.load(f)
         
-        # Load links data
-        df = pd.read_excel(LINKS_PATH)
+        # Load document links
+        data['df'] = pd.read_excel('finalcleaned_excel_file.xlsx')
         
         print("Data loaded successfully")
         return True
@@ -87,12 +82,84 @@ def load_data():
         print(f"Error loading data: {e}")
         return False
 
+def translate_text(text, source_to_target='ar_to_en'):
+    """Translate text between Arabic and English"""
+    try:
+        if source_to_target == 'ar_to_en':
+            tokenizer = models['ar_to_en_tokenizer']
+            model = models['ar_to_en_model']
+        else:
+            tokenizer = models['en_to_ar_tokenizer']
+            model = models['en_to_ar_model']
+            
+        inputs = tokenizer(text, return_tensors="pt", truncation=True)
+        outputs = model.generate(**inputs)
+        return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    except Exception as e:
+        print(f"Translation error: {e}")
+        return text
+
+def query_embeddings(query_embedding, n_results=5):
+    """Find relevant documents using embedding similarity"""
+    doc_ids = list(data['embeddings'].keys())
+    doc_embeddings = np.array(list(data['embeddings'].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]
+
+def retrieve_document_text(doc_id):
+    """Retrieve document text from HTML file"""
+    try:
+        with open(f"downloaded_articles/{doc_id}", '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 extract_entities(text):
+    """Extract medical entities from text"""
+    try:
+        results = models['ner_pipeline'](text)
+        return list({result['word'] for result in results if result['entity'].startswith("B-")})
+    except Exception as e:
+        print(f"Error extracting entities: {e}")
+        return []
+
+def generate_answer(query, context, max_length=860, temperature=0.2):
+    """Generate answer using LLM"""
+    try:
+        prompt = f"""
+        As a medical expert, answer the following question based only on the provided context:
+        
+        Context: {context}
+        Question: {query}
+        
+        Answer:"""
+        
+        inputs = models['llm_tokenizer'](prompt, return_tensors="pt", truncation=True)
+        outputs = models['llm_model'].generate(
+            inputs.input_ids,
+            max_length=max_length,
+            num_return_sequences=1,
+            temperature=temperature,
+            pad_token_id=models['llm_tokenizer'].eos_token_id
+        )
+        
+        answer = models['llm_tokenizer'].decode(outputs[0], skip_special_tokens=True)
+        return answer.split("Answer:")[-1].strip()
+    except Exception as e:
+        print(f"Error generating answer: {e}")
+        return "Sorry, I couldn't generate an answer at this time."
+
 @app.route('/health', methods=['GET'])
 def health_check():
+    """Health check endpoint"""
     return jsonify({'status': 'healthy'})
 
 @app.route('/api/query', methods=['POST'])
 def process_query():
+    """Main query processing endpoint"""
     try:
         data = request.json
         if not data or 'query' not in data:
@@ -101,24 +168,33 @@ def process_query():
         query_text = data['query']
         language_code = data.get('language_code', 0)
 
-        # Process query
+        # Translate if Arabic
         if language_code == 0:
-            query_text = translate_ar_to_en(query_text)
+            query_text = translate_text(query_text, 'ar_to_en')
 
-        # Get embeddings and find relevant documents
-        query_embedding = embedding_model.encode([query_text])
-        initial_results = query_embeddings(query_embedding, embeddings_data)
+        # Get query embedding and find relevant documents
+        query_embedding = models['embedding'].encode([query_text])
+        relevant_docs = query_embeddings(query_embedding)
+        
+        # Retrieve and process documents
+        doc_texts = [retrieve_document_text(doc_id) for doc_id, _ in relevant_docs]
         
-        # Process documents
-        document_texts = retrieve_document_texts([doc_id for doc_id, _ in initial_results])
-        relevant_portions = extract_relevant_portions(document_texts, query_text)
+        # Extract entities and generate context
+        query_entities = extract_entities(query_text)
+        contexts = []
+        for text in doc_texts:
+            doc_entities = extract_entities(text)
+            if set(query_entities) & set(doc_entities):
+                contexts.append(text)
+        
+        context = " ".join(contexts[:3])  # Use top 3 most relevant contexts
         
         # Generate answer
-        combined_text = " ".join([item for sublist in relevant_portions.values() for item in sublist])
-        answer = generate_answer(query_text, combined_text)
+        answer = generate_answer(query_text, context)
         
+        # Translate back if needed
         if language_code == 0:
-            answer = translate_en_to_ar(answer)
+            answer = translate_text(answer, 'en_to_ar')
 
         return jsonify({
             'answer': answer,
@@ -131,375 +207,13 @@ def process_query():
             'success': False
         }), 500
 
-def translate_ar_to_en(text):
-    try:
-        inputs = ar_to_en_tokenizer(text, return_tensors="pt", truncation=True)
-        outputs = ar_to_en_model.generate(**inputs)
-        return ar_to_en_tokenizer.decode(outputs[0], skip_special_tokens=True)
-    except Exception as e:
-        print(f"Translation error (AR->EN): {e}")
-        return text
-
-def translate_en_to_ar(text):
-    try:
-        inputs = en_to_ar_tokenizer(text, return_tensors="pt", truncation=True)
-        outputs = en_to_ar_model.generate(**inputs)
-        return en_to_ar_tokenizer.decode(outputs[0], skip_special_tokens=True)
-    except Exception as e:
-        print(f"Translation error (EN->AR): {e}")
-        return text
-
-language_code = 0
-
-query_text = 'How can a patient with chronic kidney disease manage their daily activities and maintain quality of life?'         #'symptoms of a heart attack '
-
-def process_query(query_text):
-    if language_code == 0:
-        # Translate Arabic input to English
-        query_text = translate_ar_to_en(query_text)
-    return query_text
-
-def embed_query_text(query_text):
-    query_embedding = embedding_model.encode([query_text])
-    return query_embedding
-
-def query_embeddings(query_embedding, embeddings_data, n_results=5):
-    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]
-    top_docs = [(doc_ids[i], similarities[i]) for i in top_indices]
-
-    return top_docs
-
-query_embedding = embed_query_text(query_text)  # Embed the query text
-initial_results = query_embeddings(query_embedding, embeddings_data, n_results=5)
-document_ids = [doc_id for doc_id, _ in initial_results]
-print(document_ids)
-
-import pandas as pd
-import requests
-from bs4 import BeautifulSoup
-
-# Load the Excel file
-file_path = '/kaggle/input/final-links/finalcleaned_excel_file.xlsx'
-df = pd.read_excel(file_path)
-
-
-# Create a dictionary mapping file names to URLs
-# Assuming the DataFrame index corresponds to file names
-file_name_to_url = {f"article_{index}.html": url for index, url in enumerate(df['Unnamed: 0'])}
-def get_page_title(url):
-    try:
-        response = requests.get(url)
-        if response.status_code == 200:
-            soup = BeautifulSoup(response.text, 'html.parser')
-            title = soup.find('title')
-            return title.get_text() if title else "No title found"
-        else:
-            return None
-    except requests.exceptions.RequestException:
-        return None
-# Example file names
-file_names = document_ids
-
-# Retrieve original URLs
-for file_name in file_names:
-    original_url = file_name_to_url.get(file_name, None)
-    if original_url:
-        title = get_page_title(original_url)
-        if title:
-            print(f"Title: {title},URL: {original_url}")
-        else:
-            print(f"Name: {file_name}")
-    else:
-        print(f"Name: {file_name}")
-
-def retrieve_document_texts(doc_ids, folder_path):
-    texts = []
-    for doc_id in doc_ids:
-        file_path = os.path.join(folder_path, doc_id)
-        try:
-            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 FileNotFoundError:
-            texts.append("")
-    return texts
-document_ids = [doc_id for doc_id, _ in initial_results]
-document_texts = retrieve_document_texts(document_ids, folder_path)
-
-# Rerank the results using the CrossEncoder
-scores = cross_encoder.predict([(query_text, doc) for doc in document_texts])
-scored_documents = list(zip(scores, document_ids, document_texts))
-scored_documents.sort(key=lambda x: x[0], reverse=True)
-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}")
-
-from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
-import nltk
-
-# Load BioBERT model and tokenizer for NER
-bio_tokenizer = AutoTokenizer.from_pretrained("blaze999/Medical-NER")
-bio_model = AutoModelForTokenClassification.from_pretrained("blaze999/Medical-NER")
-ner_biobert = pipeline("ner", model=bio_model, tokenizer=bio_tokenizer)
-
-def extract_entities(text, ner_pipeline):
-    """
-    Extract entities using a NER pipeline.
-    Args:
-        text (str): The text from which to extract entities.
-        ner_pipeline (pipeline): The NER pipeline for entity extraction.
-    Returns:
-        List[str]: A list of unique extracted entities.
-    """
-    ner_results = ner_pipeline(text)
-    entities = {result['word'] for result in ner_results if result['entity'].startswith("B-")}
-    return list(entities)
-
-def match_entities(query_entities, sentence_entities):
-    """
-    Compute the relevance score based on entity matching.
-    Args:
-        query_entities (List[str]): Entities extracted from the query.
-        sentence_entities (List[str]): Entities extracted from the sentence.
-    Returns:
-        float: The relevance score based on entity overlap.
-    """
-    query_set, sentence_set = set(query_entities), set(sentence_entities)
-    matches = query_set.intersection(sentence_set)
-    return len(matches)
-
-def extract_relevant_portions(document_texts, query, max_portions=3, portion_size=1, min_query_words=1):
-    """
-    Extract relevant text portions from documents based on entity matching.
-    Args:
-        document_texts (List[str]): List of document texts.
-        query (str): The query text.
-        max_portions (int): Maximum number of relevant portions to extract per document.
-        portion_size (int): Number of sentences to include in each portion.
-        min_query_words (int): Minimum number of matching entities to consider a sentence relevant.
-    Returns:
-        Dict[str, List[str]]: Relevant portions for each document.
-    """
-    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)
-                end_idx = min(len(sentences), i + portion_size // 2 + 1)
-                portion = " ".join(sentences[start_idx:end_idx])
-                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
-        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)
-
-        relevant_portions[f"Document_{doc_id}"] = doc_relevant_portions
-
-    return relevant_portions
-
-# Extract relevant portions based on query and documents
-relevant_portions = extract_relevant_portions(document_texts, query_text, max_portions=3, portion_size=1, min_query_words=1)
-
-for doc_id, portions in relevant_portions.items():
-    print(f"{doc_id}: {portions}")
-
-# Remove duplicates from the selected 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
-
-# Flatten the dictionary of relevant portions (from earlier code)
-flattened_relevant_portions = []
-for doc_id, portions in relevant_portions.items():
-    flattened_relevant_portions.extend(portions)
-
-# Remove duplicate portions
-unique_selected_parts = remove_duplicates(flattened_relevant_portions)
-
-# Combine the unique parts into a single string of context
-combined_parts = " ".join(unique_selected_parts)
-
-# Construct context as a list: first the query, then the unique selected portions
-context = [query_text] + unique_selected_parts
-
-# Print the context (query + relevant portions)
-print(context)
-
-import pickle
-
-with open('/kaggle/input/art-embeddings-pkl/embeddings.pkl', 'rb') as file:
-    data = pickle.load(file)
-
-# Print the type of data
-print(f"Data type: {type(data)}")
-
-# Print the first few keys and values from the dictionary
-print("First few keys and values:")
-for i, (key, value) in enumerate(data.items()):
-    if i >= 5:  # Limit to printing the first 5 key-value pairs
-        break
-    print(f"Key: {key}, Value: {value}")
-
-import pickle
-import pickletools
-
-# Load the pickle file
-file_path = '/kaggle/input/art-embeddings-pkl/embeddings.pkl'
-
-with open(file_path, 'rb') as f:
-    # Read the pickle file
-    data = pickle.load(f)
-
-# Check for suspicious or corrupted entries
-def inspect_pickle(data):
-    for key, value in data.items():
-        if isinstance(value, (str, bytes)):
-            # Try to decode and catch any non-ASCII issues
-            try:
-                value.decode('ascii')
-            except UnicodeDecodeError as e:
-                print(f"Non-ASCII entry found in key: {key}")
-                print(f"Corrupted data: {value} ({e})")
-            continue
-
-        if isinstance(value, list) and any(isinstance(v, (list, dict, str, bytes)) for v in value):
-            # Inspect list elements recursively
-            inspect_pickle({f"{key}[{idx}]": v for idx, v in enumerate(value)})
-
-# Inspect the data
-inspect_pickle(data)
-
-from transformers import AutoTokenizer, AutoModelForTokenClassification
-import torch
-import time
-
-# Load Biobert model and tokenizer
-biobert_tokenizer = AutoTokenizer.from_pretrained("blaze999/Medical-NER")
-biobert_model = AutoModelForTokenClassification.from_pretrained("blaze999/Medical-NER")
-
-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
-
-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):
-    prompt = ("""
-    As a medical expert, you are required to answer the following question based only on the provided passage. Do not include any information not present in the passage. Your response should directly reflect the content of the passage. Maintain accuracy and relevance to the provided information.
-
-    Passage: {passage}
-
-    Question: {question}
-
-    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()
-
-    output_ids = model_f.generate(
-        inputs.input_ids,
-        max_length=max_length,
-        num_return_sequences=1,
-        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())
-    answer_keywords = set(answer.lower().split())
-
-    if passage_keywords.intersection(answer_keywords):
-        return answer, duration
-    else:
-        return "Sorry, I can't help with that.", duration
-
-# Integrate Biobert model
-entities = extract_entities(query_text)
-passage = enhance_passage_with_entities(combined_parts, entities)
-# Generate answer with the enhanced passage
-prompt = create_prompt(query_text, passage)
-answer, generation_time = generate_answer(prompt)
-print(f"\nTime taken to generate the answer: {generation_time:.2f} seconds")
-def remove_answer_prefix(text):
-    prefix = "Answer:"
-    if prefix in text:
-        return text.split(prefix)[-1].strip()
-    return text
-
-def remove_incomplete_sentence(text):
-    # Check if the text ends with a period
-    if not text.endswith('.'):
-        # Find the last period or the end of the string
-        last_period_index = text.rfind('.')
-        if last_period_index != -1:
-            # Remove everything after the last period
-            return text[:last_period_index + 1].strip()
-    return text
-# Clean and print the answer
-answer_part = answer.split("Answer:")[-1].strip()
-cleaned_answer = remove_answer_prefix(answer_part)
-final_answer = remove_incomplete_sentence(cleaned_answer)
+# Initialize everything when the app starts
+print("Initializing application...")
+init_success = init_nltk() and load_models() and load_data()
 
-if language_code == 0:
-    final_answer = translate_en_to_ar(final_answer)
+if not init_success:
+    print("Failed to initialize application")
+    exit(1)
 
-if final_answer:
-    print("Answer:")
-    print(final_answer)
-else:
-    print("Sorry, I can't help with that.")
+if __name__ == "__main__":
+    app.run(host='0.0.0.0', port=7860)