initial commit

Dockerfile

@@ -0,0 +1,9 @@
+FROM python:3.11
+
+WORKDIR /app
+
+COPY . /app/
+
+RUN pip install --no-cache-dir --upgrade -r /app/requirements.txt
+
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]

README.md

@@ -1,11 +1,77 @@
----
-title: Genai
-emoji: 👁
-colorFrom: pink
-colorTo: pink
-sdk: docker
-pinned: false
-license: mit
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+# Modular RAG
+
+A hybrid approach to implement RAG inspired by Advance RAG.
+Usually implemeted with modules acting as plug and play.
+
+## Documentation
+
+#### Generator: 
+Core component of RAG, responsible for transforming the retrieved information into natural and human sense.
+
+#### Retriever: 
+The word "R" in RAG, serving the purpose of retrieving the top K element from knowledge base.
+
+#### ReRank: 
+As the name suggest a model used to re-rank the relevant documents. It indexes the documents based on the similariy score between question and the retrieved documents post vector search.
+
+## Run Locally
+
+Clone the project
+
+```bash
+  git clone https://github.com/gauravprasadgp/modular-rag
+```
+
+Go to the project directory
+
+```bash
+  cd modular-rag
+```
+
+Install dependencies
+
+```bash
+pip install -r requirements.txt
+```
+Run postgres locally
+```bash
+cd pgvector
+```
+```bash
+docker compose -d up
+```
+
+Start the server
+
+```bash
+  python main.py
+```
+
+## API Reference
+
+#### Upload file to create embedding
+
+```http
+  POST /create
+```
+
+| Parameter | Type   | Description                  |
+|:----------|:-------|:-----------------------------|
+| `file`    | `file` | **Required**. File to upload |
+
+#### Get answer from user query
+
+```http
+  POST /answer
+```
+
+| Parameter | Type     | Description              |
+|:----------| :------- |:-------------------------|
+| `query`   | `string` | **Required**. user query |
+
+
+
+## License
+
+[MIT](https://choosealicense.com/licenses/mit/)

generator/llm_calls.py

@@ -0,0 +1,33 @@
+from llama_index.legacy.embeddings import HuggingFaceEmbedding
+from llama_index.legacy.llms import LlamaCPP
+from llama_index.llms.llama_cpp.llama_utils import (
+    messages_to_prompt,
+    completion_to_prompt,
+)
+
+llm = LlamaCPP(
+    model_url="https://huggingface.co/TheBloke/Llama-2-13B-chat-GGML/resolve/main/llama-2-13b-chat.ggmlv3.q4_0"
+              ".bin",
+    temperature=0.1,
+    max_new_tokens=256,
+    context_window=3900,
+    generate_kwargs={},
+    model_kwargs={"n_gpu_layers": 1},
+    messages_to_prompt=messages_to_prompt,
+    completion_to_prompt=completion_to_prompt,
+    verbose=True,
+)
+
+
+def get_embed_model():
+    embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")
+    return embed_model
+
+
+async def get_answer(query, context):
+    prompt = f"""Given the context below answer the question.
+            Context: {context}
+            Question: {query}
+            Answer:
+            """
+    return await llm.acomplete(prompt=prompt)

main.py

@@ -0,0 +1,52 @@
+from contextlib import asynccontextmanager
+
+import uvicorn
+from fastapi import FastAPI, Request, UploadFile
+from fastapi.middleware.cors import CORSMiddleware
+
+from generator.llm_calls import get_answer
+from rerank.rerank import rerank_documents
+from retrieve.vector_store import create_embeddings_from_file, get_relevant_document
+from utils.db import postgres_db
+
+app = FastAPI(title="Modular RAG",
+              version="1.0.0", )
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    await postgres_db.create_connection_pool()
+    yield
+    await postgres_db.close_connection_pool()
+
+
+if __name__ == "__main__":
+    uvicorn.run("main:app", host="0.0.0.0", port=7860, reload=True)
+
+
+@app.post("/create")
+async def create_embedding(file: UploadFile):
+    await create_embeddings_from_file(file)
+
+
+@app.post("/answer")
+async def post_conversation(request: Request):
+    payload = await request.json()
+    query = payload.get("query")
+    context = await get_relevant_document(query=query)
+    sorted_docs = rerank_documents(question=query, documents=context)
+    sorted_context = "\n\n".join(sorted_docs)
+    return await get_answer(context=sorted_context, query=query)
+
+
+@app.get("/")
+async def get_test(request: Request):
+    return "Successfully Deployed"

pgvector/docker-compose.yml

@@ -0,0 +1,14 @@
+services:
+  db:
+    hostname: db
+    image: pgvector/pgvector:pg16
+    ports:
+     - 5432:5432
+    restart: always
+    environment:
+      - POSTGRES_DB=vectordb
+      - POSTGRES_USER=user
+      - POSTGRES_PASSWORD=password
+      - POSTGRES_HOST_AUTH_METHOD=trust
+    volumes:
+     - ./init.sql:/docker-entrypoint-initdb.d/init.sql

pgvector/init.sql

@@ -0,0 +1,9 @@
+CREATE EXTENSION IF NOT EXISTS vector;
+
+CREATE TABLE IF NOT EXISTS embeddings (
+  id SERIAL PRIMARY KEY,
+  embedding vector,
+  document text,
+  metadata jsonb
+  created_at timestamptz DEFAULT now()
+);

requirements.txt

@@ -0,0 +1,13 @@
+langchain
+psycopg[binary,pool]
+pgvector
+langchain-community
+uvicorn[standard]
+pypdf
+fastapi[all]
+python-multipart
+pydantic_settings
+llama-index
+gunicorn
+llama-index-llms-llama-cpp
+FlagEmbedding

rerank/rerank.py

@@ -0,0 +1,17 @@
+from FlagEmbedding import FlagReranker
+
+reranker = FlagReranker('BAAI/bge-reranker-large', use_fp16=True)
+
+
+def rerank_documents(question: str, documents: list[str]):
+    sentences = []
+    for doc in documents:
+        sentences.append((question, doc))
+    score = reranker.compute_score(sentences)
+    print(score)
+    sorted_elements = []
+    for score, doc in zip(score, documents):
+        elem = {score: score, doc: doc}
+        sorted_elements.append(elem)
+    sorted_docs = sorted(sorted_elements, key=lambda x: x.score, reverse=True)
+    return sorted_docs[:7]

retrieve/vector_store.py

@@ -0,0 +1,63 @@
+import datetime
+import json
+
+from fastapi import File
+from llama_index.core import SimpleDirectoryReader
+from llama_index.core.node_parser import SentenceWindowNodeParser
+from llama_index.core.schema import BaseNode
+
+from generator.llm_calls import get_embed_model
+from utils.db import postgres_db
+
+
+async def create_embeddings_from_file(file: File()):
+    documents = await SimpleDirectoryReader(input_files=[file]).aload_data()
+    node_parser = SentenceWindowNodeParser.from_defaults(
+        window_size=3,
+        window_metadata_key="window",
+        original_text_metadata_key="original_text",
+    )
+
+    nodes = node_parser.get_nodes_from_documents(documents)
+    for node in nodes:
+        embedding = await get_embed_model().aget_text_embedding(node.get_content())
+        node.embedding = embedding
+
+
+def get_values_from_nodes(nodes: list[BaseNode]):
+    values = []
+    dt = datetime.datetime.now()
+    for node in nodes:
+        value = (node.embedding, node.get_content(), json.dumps(node.metadata), dt)
+        values.append(value)
+    return values
+
+
+async def insert_documents(nodes):
+    try:
+        values = get_values_from_nodes(nodes)
+        async with postgres_db.db_pool as conn:
+            async with conn.cursor() as cur:
+                await cur.executemany("""
+                        INSERT INTO document_embedding (embedding, document, metadata, created_at)
+                        VALUES (%s, %s, %s, %s); 
+                    """, values)
+                await conn.commit()
+    except Exception as error:
+        print(f"insert document exception {error}")
+        await conn.rollback()
+
+
+async def get_relevant_document(query: str):
+    embedded_question = await get_embed_model().aget_query_embedding(query=query)
+    try:
+        async with postgres_db.db_pool as conn:
+            async with conn.cursor() as cur:
+                await cur.execute(f"""SELECT metadata -> 'window', 1 - (embedding <=> '{embedded_question}') AS 
+                cosine_similarity from document_embedding ORDER BY cosine_similarity DESC limit 10;""")
+                results = await cur.fetchall()
+                docs = [row[0] for row in results]
+                return docs
+    except Exception as error:
+        print(f"insert document exception {error}")
+        await conn.rollback()

utils/db.py

@@ -0,0 +1,38 @@
+import os
+from typing import Optional
+
+from psycopg_pool import AsyncConnectionPool
+
+
+def get_conn_str():
+    return f"""
+    dbname={os.getenv('POSTGRES_DB') or "vectordb"}
+    user={os.getenv('POSTGRES_USER') or "user"}
+    password={os.getenv('POSTGRES_PASSWORD') or "password"}
+    host={os.getenv('POSTGRES_HOST') or "localhost"}
+    port={os.getenv('POSTGRES_PORT') or "5432"}
+    """
+
+
+class PostgresDatabase:
+    def __init__(self):
+        self.db_pool = Optional[None]
+        self.conn = Optional[None]
+
+    async def create_connection_pool(self):
+        try:
+            self.conn = AsyncConnectionPool(conninfo=get_conn_str())
+            if self.conn:
+                self.db_pool = self.conn
+        except ConnectionError as error:
+            print(f"DB connection error {error}")
+
+    async def close_connection_pool(self):
+        try:
+            if self.db_pool:
+                await self.conn.close()
+        except Exception as error:
+            print(f"Error in closing db connection {error}")
+
+
+postgres_db = PostgresDatabase()