Update app.py

app.py

@@ -170,54 +170,10 @@ def translate_text(text, source_to_target='ar_to_en'):
         print(f"Translation error: {e}")
         return text
 
-def extract_entities(text):
-    """Extract medical entities from text using NER"""
-    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, please provide a clear and accurate answer to the following question based solely on the provided context.
-        
-        Context: {context}
-        
-        Question: {query}
-        
-        Answer: Let me help you with accurate information from reliable medical sources."""
-
-        inputs = models['llm_tokenizer'](prompt, return_tensors="pt", truncation=True)
-        
-        with torch.no_grad():
-            outputs = models['llm_model'].generate(
-                inputs.input_ids,
-                max_length=max_length,
-                num_return_sequences=1,
-                temperature=temperature,
-                do_sample=True,
-                top_p=0.9,
-                pad_token_id=models['llm_tokenizer'].eos_token_id
-            )
-        
-        response = models['llm_tokenizer'].decode(outputs[0], skip_special_tokens=True)
-        
-        if "Answer:" in response:
-            response = response.split("Answer:")[-1].strip()
-        
-        sentences = nltk.sent_tokenize(response)
-        if sentences:
-            return " ".join(sentences)
-        return response
-
-    except Exception as e:
-        print(f"Error generating answer: {e}")
-        return "I apologize, but I'm unable to generate an answer at this time. Please try again later."
-
+def embed_query_text(query_text):
+    query_embedding = embedding_model.encode([query_text])
+    return query_embedding
+    
 def query_embeddings(query_embedding, n_results=5):
     """Find relevant documents using embedding similarity"""
     if not data['embeddings']:
@@ -258,6 +214,160 @@ def rerank_documents(query, doc_texts):
         print(f"Error reranking documents: {e}")
         return np.zeros(len(doc_texts))
 
+def extract_entities(text):
+    """Extract medical entities from text using NER"""
+    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 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)
+
+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)
+                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
+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
+
+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
+        
+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
+
+answer_part = answer.split("Answer:")[-1].strip()
+cleaned_answer = remove_answer_prefix(answer_part)
+final_answer = remove_incomplete_sentence(cleaned_answer)
+
+
 @app.get("/")
 async def root():
     return {"message": "Welcome to the FastAPI application! Use the /health endpoint to check health, and /api/query for processing queries."}
@@ -287,8 +397,8 @@ async def chat_endpoint(chat_query: ChatQuery):
         rerank_scores = rerank_documents(query_text, doc_texts)
         ranked_texts = [text for _, text in sorted(zip(rerank_scores, doc_texts), reverse=True)]
         
-        context = " ".join(ranked_texts[:3])
-        answer = generate_answer(query_text, context)
+        context = [query_text] + unique_selected_parts
+        answer = remove_incomplete_sentence(query_text, context)
         
         return {
             "response": answer,