This repository implements a Retrieval-Augmented Generation (RAG) workflow that:
- Fetches abstracts from PubMed (via NCBI E-utilities)
- leans and embeds them using Sentence Transformers
- Stores embeddings in ChromaDB for efficient similarity search
- Queries domain knowledge to assist in filling missing predicates in biomedical knowledge graphs
- Uses AWS Bedrock LLMs with retrieved context to generate reasoning-based explanations
We detail the RAG usage below:
The pipeline automatically:
- Retrieves PubMed abstracts (default: Alzheimer's disease)
- Cleans, chunks, embeds, and stores them in ChromaDB
import requests
from sentence_transformers import SentenceTransformer
from langchain.text_splitter import RecursiveCharacterTextSplitter
import chromadb
# Initialize components
client = chromadb.Client()
collection = client.create_collection("pubtator_data")
model = SentenceTransformer("all-MiniLM-L6-v2")results = query_domain_knowledge(
"amyloid beta and alzheimers relationship",
collection,
model,
top_k=3
)
print(results)filled_df, metrics, responses = fill_missing_predicates_llm_with_domain_knowledge(
input_df=my_triples_df,
unique_predicates=["biolink:related_to", "biolink:interacts_with", "biolink:causes"],
collection=collection,
model=model,
output_file='bedrock_rag_filled_test.csv',
metrics_file='bedrock_rag_metrics_test.json',
responses_file='bedrock_rag_responses_test.json'
)This produces:
bedrock_rag_filled_test.csv– triples with filled predicatesbedrock_rag_metrics_test.json– run statisticsbedrock_rag_responses_test.json– detailed LLM responses