app.py

import os
import pickle
import numpy as np
from flask import Flask, request, jsonify
from flask_cors import CORS
from transformers import (
    AutoTokenizer, 
    AutoModelForSeq2SeqLM, 
    AutoModelForCausalLM,
    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
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()

# 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

        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 initializing models: {e}")
        return False

# Load data with error handling
def load_data():
    try:
        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

@app.route('/health', methods=['GET'])
def health_check():
    return jsonify({'status': 'healthy'})

@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

        query_text = data['query']
        language_code = data.get('language_code', 0)

        # Process query
        if language_code == 0:
            query_text = translate_ar_to_en(query_text)

        # Get embeddings and find relevant documents
        query_embedding = embedding_model.encode([query_text])
        initial_results = query_embeddings(query_embedding, embeddings_data)
        
        # Process documents
        document_texts = retrieve_document_texts([doc_id for doc_id, _ in initial_results])
        relevant_portions = extract_relevant_portions(document_texts, query_text)
        
        # 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)

        return jsonify({
            'answer': answer,
            'success': True
        })

    except Exception as e:
        return jsonify({
            'error': str(e),
            '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)

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.")