Remove Complexity from Your RAG Applications

def get_pdf_embedding(pdf_path: str, model, processor):
    page_images, page_texts = get_pdf_images(pdf_path=pdf_path)
    page_embeddings = []
    dataloader = DataLoader(
        page_images,
        batch_size=2,
        shuffle=False,
        collate_fn=lambda x: process_images(processor, x),
    )

    for batch_doc in tqdm(dataloader):
        with torch.no_grad():
            batch_doc = {k: v.to(model.device) for k, v in batch_doc.items()}
            embeddings_doc = model(**batch_doc)
            page_embeddings.extend(list(torch.unbind(embeddings_doc.to("cpu"))))
    document = {
        "name": pdf_path,
        "page_images": page_images,
        "page_texts": page_texts,
        "page_embeddings": page_embeddings,
    }
    return document
 

def get_query_embedding(query: str, model, processor):
    dummy_image = PIL.Image.new("RGB", (448, 448), (255, 255, 255))
    dataloader = DataLoader(
        [query],
        batch_size=4,
        shuffle=False,
        collate_fn=lambda x: process_queries(processor, x, dummy_image),
    )
    qs = []
    for batch_query in dataloader:
        with torch.no_grad():
            batch_query = {k: v.to(model.device) for k, v in batch_query.items()}
            embeddings_query = model(**batch_query)
        qs.extend(list(torch.unbind(embeddings_query.to("cpu"))))
    q = {"query": query, "embeddings": qs[0]}
    return q

def create_db(docs_storage, table_name: str = "demo", db_path: str = "lancedb"):
    db = lancedb.connect(db_path)
    data = []
    for x in docs_storage:
        sample = {
            "name": x["name"],
            "page_texts": x["page_text"],
            "image": get_base64_image(x["page_image"]),
            "page_idx": x["page_idx"],
            "page_embedding_flatten": x["page_embedding"].float().numpy().flatten(),
            "page_embedding_shape": x["page_embedding"].float().numpy().shape
        }
        data.append(sample)

    table = db.create_table(table_name, data, mode="overwrite")
    return table
    

def search(query: str, table_name: str, model, processor, db_path: str = "lancedb", top_k: int = 3):
    qs = get_query_embedding(query=query, model=model, processor=processor)
    db = lancedb.connect(db_path)
    table = db.open_table(table_name)
    # Search over all dataset
    r = table.search().limit(None).to_list()
    
    def process_patch_embeddings(x):
        patches = np.reshape(x['page_embedding_flatten'], x['page_embedding_shape'])
        return torch.from_numpy(patches).to(torch.bfloat16)
    
    all_pages_embeddings = [process_patch_embeddings(x) for x in r]
    
    retriever_evaluator = CustomEvaluator(is_multi_vector=True)
    scores = retriever_evaluator.evaluate_colbert([qs["embeddings"]], all_pages_embeddings)

    top_k_indices = torch.topk(scores, k=top_k, dim=1).indices

    results = []
    for idx in top_k_indices[0]:
        page = r[idx]
        pil_image = base64_to_pil(page["image"])
        result = {"name": page["name"], "page_idx": page["page_idx"], "pil_image": pil_image}
        results.append(result)
    return results

def run_vision_inference(input_images: PIL.Image, prompt: str, model, processor):
    images = []
    placeholder = ""

    # Note: if OOM, you might consider reduce number of frames in this example.
    for i in range(len(input_images)):
        images.append(input_images[i])
        placeholder += f"<|image_{i + 1}|>\n"

    messages = [
        {"role": "user", "content": f"{placeholder} {prompt}"},
    ]

    prompt = processor.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

    inputs = processor(prompt, images, return_tensors="pt").to("cuda:0")

    generation_args = {
        "max_new_tokens": 512,
        "temperature": 0.2,
        "do_sample": True,
    }

    generate_ids = model.generate(**inputs, eos_token_id=processor.tokenizer.eos_token_id, **generation_args)
    # remove input tokens
    generate_ids = generate_ids[:, inputs["input_ids"].shape[1] :]
    response = processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
    return response

@app.cls(
    gpu="a10g",
    ....
)
class VisionRAG:
    @modal.enter()
    def load(self):
        from vision_retrieval.core import get_model_colpali, get_model_phi_vision
        model_colpali, processor_colpali = get_model_colpali(base_model_id="/model-paligemma-3b-mix-448")
        model_phi_vision, processor_phi_vision = get_model_phi_vision(model_id="/model-phi-3.5-vision-instruct")
        self.model_colpali = model_colpali
        self.processor_colpali = processor_colpali
        self.model_phi_vision = model_phi_vision
        self.processor_phi_vision = processor_phi_vision

    @modal.method()
    def ingest_data(self, pdf_urls: List[str], table_name: str, db_path: str):
    
        print("1. Downloads PDFs")
        from tqdm import tqdm
        from vision_retrieval.core import create_db, download_pdf, embedd_docs
        pdfs = []
        for pdf_url in tqdm(pdf_urls):
            pdf_file_name = download_pdf(url=pdf_url)
            pdfs.append(pdf_file_name)
        print(f"result pdfs = {pdfs}")
        
        print("2. Generating embeddings")
        docs_storage = embedd_docs(docs_path=pdfs, model=self.model_colpali, processor=self.processor_colpali)

        print(f"result docs = {len(docs_storage)}")

        print("3. Build vectorDB")
        create_db(docs_storage=docs_storage, table_name=table_name, db_path=db_path)
        print("Done!")

    @modal.method()
    def query_data(self, user_query: str, table_name: str, db_path: str):
        from vision_retrieval.core import run_vision_inference, search

        print("1. Search relevant images")
        search_results = search(
            query=user_query,
            table_name=table_name,
            db_path=db_path,
            processor=self.processor_colpali,
            model=self.model_colpali,
        )
        print(f"result most relevant {search_results[0]}")

        print("2. Build prompt")
        prompt = f"""
        Below is a user query, I want you to answer the query using images provided.
        user query:
        {user_query}
        """

        print("3. Query LLM with prompt and relavent images")
        response = run_vision_inference(
            input_images=[search_results[0]['pil_image']], prompt=prompt, model=self.model_phi_vision, processor=self.processor_phi_vision
        )
        print(f"response = {response}")
        return {"response": response, "page": search_results[0]['page_idx'] + 1, "pdf_name": search_results[0]['name']}

import lancedb
import modal
from dagster import AssetExecutionContext, MetadataValue, asset

.... 

@asset(group_name="ingest", compute_kind="python")
def pdf_corpus(context: AssetExecutionContext):
    pdf_urls = [
        "https://vision-retrieval.s3.amazonaws.com/docs/InfraRedReport.pdf",
    ]
    context.add_output_metadata(
        {
            "len": MetadataValue.int(len(pdf_urls)),
            "sample": MetadataValue.json(pdf_urls),
        }
    )
    return pdf_urls


@asset(group_name="ingest", compute_kind="modal-lab")
def pdf_embeddings_table(context: AssetExecutionContext, pdf_corpus):
    VisionRAG = modal.Cls.lookup("vision_retrieval", "VisionRAG")
    vision_rag = VisionRAG()

    db_path = "s3://vision-retrieval/storage"
    table_name = "dagster-table"

    vision_rag.ingest_data.remote(pdf_urls=pdf_corpus, table_name=table_name, db_path=db_path)

    db = lancedb.connect(db_path)
    table_names = db.table_names()
    t = db.open_table(table_name)
    schema = t.schema

    context.add_output_metadata(
        {
            "table_names": MetadataValue.json(table_names),
            "table_schema": MetadataValue.md(schema.to_string()),
        }
    )
    return table_name


@asset(group_name="query", compute_kind="python")
def query_samples(context: AssetExecutionContext):
    query_samples_ = ["How does inference costs change over time?", "Top companies in observability space?"]
    context.add_output_metadata(
        {
            "len": MetadataValue.int(len(query_samples_)),
            "sample": MetadataValue.json(query_samples_),
        }
    )
    return query_samples_


@asset(group_name="query", compute_kind="modal-lab")
def query_results(context: AssetExecutionContext, query_samples, pdf_embeddings_table):
    VisionRAG = modal.Cls.lookup("vision_retrieval", "VisionRAG")
    vision_rag = VisionRAG()

    # TODO: move to configs
    db_path = "s3://vision-retrieval/storage"
    
    data = []
    for q in query_samples:
        result = vision_rag.query_data.remote(user_query=q, table_name=pdf_embeddings_table, db_path=db_path)
        result['query'] = q
        data.append(result)

    context.add_output_metadata(
        {
            "data": MetadataValue.json(data),
        }
    )