Update app.py

app.py

@@ -1,226 +1,505 @@
-from fastapi import FastAPI, HTTPException, Query
-from pydantic import BaseModel
-from typing import List, Optional, Dict
+import os
 import pickle
 import numpy as np
-from sklearn.metrics.pairwise import cosine_similarity
-from sentence_transformers import SentenceTransformer, CrossEncoder, util
-from bs4 import BeautifulSoup
-import os
-import nltk
-import torch
+from flask import Flask, request, jsonify
+from flask_cors import CORS
 from transformers import (
     AutoTokenizer, 
-    BartForConditionalGeneration,
+    AutoModelForSeq2SeqLM, 
     AutoModelForCausalLM,
-    AutoModelForSeq2SeqLM,
     AutoModelForTokenClassification
 )
+from sentence_transformers import SentenceTransformer, CrossEncoder
+from sklearn.metrics.pairwise import cosine_similarity
+from bs4 import BeautifulSoup
+import nltk
+import torch
 import pandas as pd
-import time
+from startup import setup_files
 
-app = FastAPI()
-
-# ArticleEmbeddingUnpickler and safe_load_embeddings functions
-class ArticleEmbeddingUnpickler(pickle.Unpickler):
-    """Custom unpickler for article embeddings with enhanced persistence handling"""
-    def find_class(self, module: str, name: str) -> any:
-        if module == 'numpy':
-            return getattr(np, name)
-        if module == 'sentence_transformers.SentenceTransformer':
-            from sentence_transformers import SentenceTransformer
-            return SentenceTransformer
-        return super().find_class(module, name)
-    
-    def persistent_load(self, pid: any) -> str:
-        """Enhanced persistent ID handler with better encoding management"""
-        try:
-            if isinstance(pid, bytes):
-                return pid.decode('utf-8', errors='replace')
-            if isinstance(pid, (str, int, float)):
-                return str(pid)
-            return repr(pid)
-        except Exception as e:
-            print(f"Warning: Error in persistent_load: {str(e)}")
-            return repr(pid)
-
-def safe_load_embeddings(file_path: str = 'embeddings.pkl') -> Dict[str, np.ndarray]:
-    """Load embeddings with enhanced error handling, validation, and persistent ID support."""
-    def persistent_load(pid):
-        print(f"Warning: Persistent ID encountered: {pid}")
-        raise ValueError("Persistent IDs are not supported in this application")
 
-    try:
-        if not os.path.exists(file_path):
-            raise FileNotFoundError(f"Embeddings file not found at {file_path}")
+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')
 
-        with open(file_path, 'rb') as file:
-            unpickler = pickle.Unpickler(file)
-            unpickler.persistent_load = persistent_load  # Assign custom handler
-            embeddings_data = unpickler.load()
+def init_app():
+    # Download and extract files if they don't exist
+    if not os.path.exists('downloaded_articles'):
+        setup_files()
 
-        if not isinstance(embeddings_data, dict):
-            raise ValueError(f"Invalid data structure: expected dict, got {type(embeddings_data)}")
+# Initialize models with proper error handling
+def initialize_models():
+    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
 
-        # Process and validate embeddings
-        valid_embeddings = {}
-        for key, value in embeddings_data.items():
-            try:
-                key_str = str(key).encode('ascii', errors='replace').decode('ascii').strip()
-                if not key_str:
-                    print(f"Skipping invalid or empty key: {key}")
-                    continue
-
-                if isinstance(value, list):
-                    value = np.array(value, dtype=np.float32)
-                elif isinstance(value, np.ndarray):
-                    value = value.astype(np.float32)
-                else:
-                    print(f"Skipping invalid embedding type for key {key_str}: {type(value)}")
-                    continue
-
-                if value.ndim != 1 or np.isnan(value).any() or np.isinf(value).any():
-                    print(f"Skipping embedding with invalid shape or values for key {key_str}")
-                    continue
-
-                valid_embeddings[key_str] = value
-            except Exception as e:
-                print(f"Error processing embedding for key {key}: {str(e)}")
-                continue
-
-        if not valid_embeddings:
-            raise ValueError("No valid embeddings found in file")
-
-        print(f"Successfully loaded {len(valid_embeddings)} valid embeddings")
-        return valid_embeddings
+        print("Initializing models...")
+        
+        # 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")
+
+        # 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)
+        
+        print("Models initialized successfully")
+        return True
     except Exception as e:
-        print(f"Error loading embeddings: {str(e)}")
-        raise
-
-def safe_save_embeddings(embeddings_dict, file_path='embeddings.pkl'):
-    cleaned_embeddings = {
-        str(key): value
-        for key, value in embeddings_dict.items()
-    }
-    with open(file_path, 'wb') as f:
-        pickle.dump(cleaned_embeddings, f, protocol=4)
-
-
-# GlobalModels and load_models
-class GlobalModels:
-    embedding_model = None
-    cross_encoder = None
-    semantic_model = None
-    tokenizer = None
-    model = None
-    tokenizer_f = None
-    model_f = None
-    ar_to_en_tokenizer = None
-    ar_to_en_model = None
-    en_to_ar_tokenizer = None
-    en_to_ar_model = None
-    bio_tokenizer = None
-    bio_model = None
-    embeddings_data = None
-    file_name_to_url = None
-
-global_models = GlobalModels()
-
-@app.on_event("startup")
-async def load_models():
+        print(f"Error initializing models: {e}")
+        return False
+
+# Load data with error handling
+def load_data():
     try:
-        global_models.embedding_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
-        global_models.cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2', max_length=512)
-        global_models.semantic_model = SentenceTransformer('all-MiniLM-L6-v2')
+        global embeddings_data, df
+        
+        print("Loading data files...")
+        
+        # Load embeddings
+        with open(EMBEDDINGS_PATH, 'rb') as file:
+            embeddings_data = pickle.load(file)
+        
+        # Load links data
+        df = pd.read_excel(LINKS_PATH)
+        
+        print("Data loaded successfully")
+        return True
+    except Exception as e:
+        print(f"Error loading data: {e}")
+        return False
 
-        global_models.tokenizer = AutoTokenizer.from_pretrained("facebook/bart-base")
-        global_models.model = BartForConditionalGeneration.from_pretrained("facebook/bart-base")
+@app.route('/health', methods=['GET'])
+def health_check():
+    return jsonify({'status': 'healthy'})
 
-        model_name = "M4-ai/Orca-2.0-Tau-1.8B"
-        global_models.tokenizer_f = AutoTokenizer.from_pretrained(model_name)
-        global_models.model_f = AutoModelForCausalLM.from_pretrained(model_name)
+@app.route('/api/query', methods=['POST'])
+def process_query():
+    try:
+        data = request.json
+        if not data or 'query' not in data:
+            return jsonify({'error': 'No query provided', 'success': False}), 400
 
-        global_models.ar_to_en_tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-ar-en")
-        global_models.ar_to_en_model = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-ar-en")
-        global_models.en_to_ar_tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-ar")
-        global_models.en_to_ar_model = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-en-ar")
+        query_text = data['query']
+        language_code = data.get('language_code', 0)
 
-        global_models.bio_tokenizer = AutoTokenizer.from_pretrained("blaze999/Medical-NER")
-        global_models.bio_model = AutoModelForTokenClassification.from_pretrained("blaze999/Medical-NER")
+        # Process query
+        if language_code == 0:
+            query_text = translate_ar_to_en(query_text)
 
-        try:
-            with open('embeddings.pkl', 'rb') as file:
-                global_models.embeddings_data = pickle.load(file)
-        except (FileNotFoundError, pickle.UnpicklingError) as e:
-            print(f"Error loading embeddings data: {e}")
-            raise HTTPException(status_code=500, detail="Failed to load embeddings data.")
+        # Get embeddings and find relevant documents
+        query_embedding = embedding_model.encode([query_text])
+        initial_results = query_embeddings(query_embedding, embeddings_data)
         
-        df = pd.read_excel('finalcleaned_excel_file.xlsx')
-        global_models.file_name_to_url = {f"article_{index}.html": url for index, url in enumerate(df['Unnamed: 0'])}
+        # Process documents
+        document_texts = retrieve_document_texts([doc_id for doc_id, _ in initial_results])
+        relevant_portions = extract_relevant_portions(document_texts, query_text)
         
-    except Exception as e:
-        print(f"Error loading models: {e}")
-        raise HTTPException(status_code=500, detail="Failed to load models.")
-
+        # Generate answer
+        combined_text = " ".join([item for sublist in relevant_portions.values() for item in sublist])
+        answer = generate_answer(query_text, combined_text)
+        
+        if language_code == 0:
+            answer = translate_en_to_ar(answer)
 
-# Query and Document Retrieval Endpoint
-class QueryInput(BaseModel):
-    query_text: str
-    language_code: int  # 0 for Arabic, 1 for English
-    query_type: str    # "profile" or "question"
-    previous_qa: Optional[List[Dict[str, str]]] = None
+        return jsonify({
+            'answer': answer,
+            'success': True
+        })
 
-class DocumentResponse(BaseModel):
-    title: str
-    url: str
-    text: str
-    score: float
+    except Exception as e:
+        return jsonify({
+            'error': str(e),
+            'success': False
+        }), 500
 
-@app.post("/retrieve_documents")
-async def retrieve_documents(input_data: QueryInput):
+def translate_ar_to_en(text):
     try:
-        processed_query = process_query(input_data.query_text, input_data.language_code)
-        query_embedding = embed_query_text(processed_query)
-        results = query_embeddings(query_embedding)
-        
-        document_ids = [doc_id for doc_id, _ in results]
-        document_texts = retrieve_document_texts(document_ids)
-        scores = global_models.cross_encoder.predict([(processed_query, doc) for doc in document_texts])
-        
-        documents = []
-        for score, doc_id, text in zip(scores, document_ids, document_texts):
-            url = global_models.file_name_to_url.get(doc_id, "")
-            documents.append({
-                "title": doc_id,
-                "url": url,
-                "text": text if input_data.language_code == 1 else translate_en_to_ar(text),
-                "score": float(score)
-            })
-        return documents
+        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:
-        raise HTTPException(status_code=500, detail=f"Error retrieving documents: {str(e)}")
+        print(f"Translation error (EN->AR): {e}")
+        return text
+
+language_code = 0
 
-def process_query(query_text: str, language_code: int) -> str:
+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:
-        return translate_ar_to_en(query_text)  # Translate Arabic to English if required
+        # Translate Arabic input to English
+        query_text = translate_ar_to_en(query_text)
     return query_text
 
-def embed_query_text(query_text: str) -> np.ndarray:
-    return global_models.embedding_model.encode(query_text, convert_to_tensor=True)
+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)
 
-def query_embeddings(query_embedding: np.ndarray, top_n: int = 10) -> List[tuple]:
-    doc_embeddings = list(global_models.embeddings_data.values())
-    document_ids = list(global_models.embeddings_data.keys())
-    similarities = cosine_similarity(query_embedding, doc_embeddings)
-    top_indices = np.argsort(similarities[0])[-top_n:]
-    return [(document_ids[idx], similarities[0][idx]) for idx in reversed(top_indices)]
+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
 
-def retrieve_document_texts(document_ids: List[str]) -> List[str]:
-    return [global_models.file_name_to_url[doc_id] for doc_id in document_ids]
+# 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)
 
-def translate_en_to_ar(text: str) -> str:
-    # Translation logic here, possibly using `transformers` or another library
-    pass
+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)
+
+if language_code == 0:
+    final_answer = translate_en_to_ar(final_answer)
+
+if final_answer:
+    print("Answer:")
+    print(final_answer)
+else:
+    print("Sorry, I can't help with that.")