Create rag_code.py

rag_code.py

@@ -0,0 +1,211 @@
+import torch
+from qdrant_client import models
+from qdrant_client import QdrantClient
+from colpali_engine.models import ColPali, ColPaliProcessor
+from Janus.janus.models import MultiModalityCausalLM, VLChatProcessor
+from Janus.janus.utils.io import load_pil_images
+from transformers import AutoModelForCausalLM
+import base64
+from io import BytesIO
+from tqdm import tqdm
+
+def batch_iterate(lst, batch_size):
+    """Yield successive n-sized chunks from lst."""
+    for i in range(0, len(lst), batch_size):
+        yield lst[i : i + batch_size]
+
+def image_to_base64(image):
+    buffered = BytesIO()
+
+    image.save(buffered, format="JPEG")
+
+    return base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+class EmbedData:
+
+    def __init__(self, embed_model_name="vidore/colpali-v1.2", batch_size = 4):
+        self.embed_model_name = embed_model_name
+        self.embed_model, self.processor = self._load_embed_model()
+        self.batch_size = batch_size
+        self.embeddings = []
+        
+    def _load_embed_model(self):
+        embed_model = ColPali.from_pretrained(
+            self.embed_model_name,
+            torch_dtype=torch.bfloat16,
+            device_map="mps",
+            trust_remote_code=True, 
+            cache_dir="./Janus/hf_cache"
+        )
+
+        processor = ColPaliProcessor.from_pretrained(self.embed_model_name)
+        return embed_model, processor
+    
+    def get_query_embedding(self, query):
+        
+        with torch.no_grad():
+            query = self.processor.process_queries([query]).to(self.embed_model.device)
+            
+            query_embedding = self.embed_model(**query)
+
+        return query_embedding[0].cpu().float().numpy().tolist()
+
+    def generate_embedding(self, images):
+        with torch.no_grad():
+            batch_images = self.processor.process_images(images).to(self.embed_model.device)
+            image_embeddings = self.embed_model(**batch_images).cpu().float().numpy().tolist()
+        
+        return image_embeddings
+        
+    def embed(self, images):
+        
+        self.images = images
+        self.all_embeddings = []
+
+        for batch_images in tqdm(batch_iterate(images, self.batch_size), desc="Generating embeddings"):
+            batch_embeddings = self.generate_embedding(batch_images)
+            self.embeddings.extend(batch_embeddings)
+
+class QdrantVDB_QB:
+
+    def __init__(self, collection_name, vector_dim = 128, batch_size=4):
+        self.collection_name = collection_name
+        self.batch_size = batch_size
+        self.vector_dim = vector_dim
+        
+    def define_client(self):
+        
+        self.client = QdrantClient(url="http://localhost:6333", prefer_grpc=True)
+        
+    def create_collection(self):
+        
+        if not self.client.collection_exists(collection_name=self.collection_name):
+
+            self.client.create_collection(
+                collection_name=self.collection_name,
+                on_disk_payload=True,
+                vectors_config=models.VectorParams(
+                    size=self.vector_dim,
+                    distance=models.Distance.COSINE,
+                    on_disk=True,
+                    multivector_config=models.MultiVectorConfig(
+                        comparator=models.MultiVectorComparator.MAX_SIM
+                    ),
+                ),
+            )
+
+    def ingest_data(self, embeddata):
+    
+        for i, batch_embeddings in tqdm(enumerate(batch_iterate(embeddata.embeddings, self.batch_size)), desc="Ingesting data"):
+
+            points = []
+            for j, embedding in enumerate(batch_embeddings):
+
+                image_bs64 = image_to_base64(embeddata.images[i*self.batch_size + j])
+                                    
+                current_point = models.PointStruct(id=i*self.batch_size + j,
+                                                   vector=embedding,
+                                                   payload={"image": image_bs64})
+
+                points.append(current_point)
+    
+            self.client.upsert(collection_name=self.collection_name, points=points, wait=True)
+        
+class Retriever:
+
+    def __init__(self, vector_db, embeddata):
+        
+        self.vector_db = vector_db
+        self.embeddata = embeddata
+
+    def search(self, query):
+        query_embedding = self.embeddata.get_query_embedding(query)
+        
+        query_result = self.vector_db.client.query_points(collection_name=self.vector_db.collection_name,
+                                        query=query_embedding,
+                                        limit=4,
+                                        search_params=models.SearchParams(
+                                        quantization=models.QuantizationSearchParams(
+                                        ignore=True,
+                                            rescore=True,
+                                            oversampling=2.0
+                                            )
+                                        )
+                                    )
+
+        return query_result
+    
+class RAG:
+
+    def __init__(self,
+                 retriever,
+                 llm_name = "deepseek-ai/Janus-Pro-1B"
+                 ):
+        
+        self.llm_name = llm_name
+        self._setup_llm()
+        self.retriever = retriever
+
+    def _setup_llm(self):
+
+        self.vl_chat_processor = VLChatProcessor.from_pretrained(self.llm_name, cache_dir="./Janus/hf_cache")
+        self.tokenizer = self.vl_chat_processor.tokenizer
+
+        self.vl_gpt = AutoModelForCausalLM.from_pretrained(
+            self.llm_name, trust_remote_code=True, cache_dir="./Janus/hf_cache"
+        ).to(torch.bfloat16).eval()
+
+    def generate_context(self, query):
+
+        result = self.retriever.search(query)
+        return f"./images/page{result.points[0].id}.jpg"
+
+    def query(self, query):
+        image_context = self.generate_context(query=query)
+
+        qa_prompt_tmpl_str = f"""The user has asked the following question:
+
+                        ---------------------
+                        
+                        Query: {query}
+                        
+                        ---------------------
+
+                        Some images are available to you
+                        for this question. You have
+                        to understand these images thoroughly and 
+                        extract all relevant information that will 
+                        help you answer the query.
+                                     
+                        ---------------------
+                        """
+        
+        conversation = [
+            {
+                "role": "User",
+                "content": f"<image_placeholder> \n {qa_prompt_tmpl_str}",
+                "images": [image_context],
+            },
+            {"role": "Assistant", "content": ""},
+        ]
+
+        pil_images = load_pil_images(conversation)
+        prepare_inputs = self.vl_chat_processor(
+            conversations=conversation, images=pil_images, force_batchify=True
+        ).to(self.vl_gpt.device)
+
+        inputs_embeds = self.vl_gpt.prepare_inputs_embeds(**prepare_inputs)
+
+        outputs = self.vl_gpt.language_model.generate(
+            inputs_embeds=inputs_embeds,
+            attention_mask=prepare_inputs.attention_mask,
+            pad_token_id=self.tokenizer.eos_token_id,
+            bos_token_id=self.tokenizer.bos_token_id,
+            eos_token_id=self.tokenizer.eos_token_id,
+            max_new_tokens=512,
+            do_sample=False,
+            use_cache=True,
+        )
+        streaming_response = self.tokenizer.decode(outputs[0].cpu().tolist(), skip_special_tokens=True)
+        
+        return streaming_response