started RAG

.gitignore

@@ -1,2 +1,3 @@
 usage.csv
 temp_95d8748770ec4615942cc78d473f5cee.csv
+uploads/*

client.py

@@ -5,6 +5,17 @@ import streamlit as st
 # stui = StreamlitUI(api_url="http://localhost:8000")  # FastAPI backend URL
 st.title("FastAPI ChatBot") 
 
+st.write("Upload a file to FastAPI")
+file = st.file_uploader("Choose a file", type=["pdf"])
+
+if st.button("Submit"):
+    if file is not None:
+        files = {"file": (file.name, file, file.type)}
+        response = requests.post("http://localhost:8000/upload", files=files)
+        st.write(response.text)
+    else:
+        st.write("No file uploaded.")
+        
 if "messages" not in st.session_state:
     st.session_state.messages = [] 
 

main.py

@@ -1,11 +1,11 @@
 # main.py
-from fastapi import FastAPI,status,Response,Request,Depends
+from fastapi import FastAPI,status,Response,Request,Depends,HTTPException,UploadFile, File
 from fastapi.responses import StreamingResponse,FileResponse
 from models import load_text_model,generate_text,load_audio_model,generate_audio,load_image_model, generate_image
 from schemas import VoicePresets
 from utils import audio_array_to_buffer,img_to_bytes
 from contextlib import asynccontextmanager
-from typing import AsyncIterator,Callable,Awaitable
+from typing import AsyncIterator,Callable,Awaitable,Annotated
 from uuid import uuid4
 import time
 from datetime import datetime, timezone
@@ -13,13 +13,14 @@ import csv
 from dependencies import get_urls_content
 from schemas import TextModelResponse,TextModelRequest
 import shutil, uuid
+from upload import save_file
 
 models = {} 
 
 @asynccontextmanager 
 async def lifespan(_: FastAPI) -> AsyncIterator[None]:
     # models["text2image"] = load_image_model() 
-    models["text"]=load_text_model()
+    # models["text"]=load_text_model()
     yield 
     models.clear() 
 
@@ -112,4 +113,22 @@ def serve_text_to_image_model_controller(prompt: str):
     # pipe = load_image_model()
     # output = generate_image(pipe, prompt) 
     output = generate_image(models["text2image"], prompt)
-    return Response(content=img_to_bytes(output), media_type="image/png") 
+    return Response(content=img_to_bytes(output), media_type="image/png") 
+
+@app.post("/upload")
+async def file_upload_controller(
+    file: Annotated[UploadFile, File(description="Uploaded PDF documents")]
+):
+    if file.content_type != "application/pdf":
+        raise HTTPException(
+            detail=f"Only uploading PDF documents are supported",
+            status_code=status.HTTP_400_BAD_REQUEST,
+        )
+    try:
+        await save_file(file)
+    except Exception as e:
+        raise HTTPException(
+            detail=f"An error occurred while saving file - Error: {e}",
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+        )
+    return {"filename": file.filename, "message": "File uploaded successfully"}

rag/extractor.py

@@ -0,0 +1,11 @@
+from pypdf import PdfReader
+
+def pdf_text_extractor(filepath: str) -> None:
+    content = ""
+    pdf_reader = PdfReader(filepath, strict=True) 
+    for page in pdf_reader.pages:
+        page_text = page.extract_text()
+        if page_text:
+            content += f"{page_text}\n\n" 
+    with open(filepath.replace("pdf", "txt"), "w", encoding="utf-8") as file: 
+        file.write(content)

rag/transform.py

@@ -0,0 +1,28 @@
+import re
+from typing import Any, AsyncGenerator
+
+import aiofiles
+from transformers import AutoModel
+
+DEFAULT_CHUNK_SIZE = 1024 * 1024 * 50  # 50 megabytes
+
+embedder = AutoModel.from_pretrained(
+    "jinaai/jina-embeddings-v2-base-en", trust_remote_code=True 
+)
+
+async def load(filepath: str) -> AsyncGenerator[str, Any]:
+    async with aiofiles.open(filepath, "r", encoding="utf-8") as f: 
+        while chunk := await f.read(DEFAULT_CHUNK_SIZE): 
+            yield chunk 
+
+def clean(text: str) -> str:
+    t = text.replace("\n", " ")
+    t = re.sub(r"\s+", " ", t)
+    t = re.sub(r"\. ,", "", t)
+    t = t.replace("..", ".")
+    t = t.replace(". .", ".")
+    cleaned_text = t.replace("\n", " ").strip()
+    return cleaned_text 
+
+def embed(text: str) -> list[float]:
+    return embedder.encode(text).tolist() 

requirements.txt

@@ -10,4 +10,8 @@ loguru
 beautifulsoup4 
 lxml 
 aiohttp
-tiktoken
+tiktoken
+aiofiles 
+python-multipart
+qdrant_client 
+pypdf 

upload.py

@@ -0,0 +1,14 @@
+import os
+import aiofiles
+from aiofiles.os import makedirs
+from fastapi import UploadFile
+
+DEFAULT_CHUNK_SIZE = 1024 * 1024 * 50  # 50 megabytes
+
+async def save_file(file: UploadFile) -> str:
+    await makedirs("uploads", exist_ok=True)
+    filepath = os.path.join("uploads", file.filename)
+    async with aiofiles.open(filepath, "wb") as f:
+        while chunk := await file.read(DEFAULT_CHUNK_SIZE):
+            await f.write(chunk)
+    return filepath