initial commit
This commit is contained in:
13
.gitignore
vendored
Normal file
13
.gitignore
vendored
Normal file
@@ -0,0 +1,13 @@
|
||||
# Python-generated files
|
||||
__pycache__/
|
||||
*.py[oc]
|
||||
build/
|
||||
dist/
|
||||
wheels/
|
||||
*.egg-info
|
||||
|
||||
# Virtual environments
|
||||
.venv
|
||||
|
||||
|
||||
.env
|
||||
83
main.py
Normal file
83
main.py
Normal file
@@ -0,0 +1,83 @@
|
||||
import ollama
|
||||
from uuid import uuid4, UUID
|
||||
|
||||
from request import PaperlessNGXService
|
||||
|
||||
from math import ceil
|
||||
|
||||
import chromadb
|
||||
|
||||
client = chromadb.EphemeralClient()
|
||||
collection = client.create_collection(name="docs")
|
||||
|
||||
|
||||
class Chunk:
|
||||
def __init__(
|
||||
self,
|
||||
text: str,
|
||||
size: int,
|
||||
document_id: UUID,
|
||||
chunk_id: int,
|
||||
embedding,
|
||||
):
|
||||
self.text = text
|
||||
self.size = size
|
||||
self.document_id = document_id
|
||||
self.chunk_id = chunk_id
|
||||
self.embedding = embedding
|
||||
|
||||
|
||||
class Chunker:
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
|
||||
def chunk_document(self, document: str, chunk_size: int = 300) -> list[Chunk]:
|
||||
doc_uuid = uuid4()
|
||||
|
||||
chunks = []
|
||||
num_chunks = ceil(len(document) / chunk_size)
|
||||
document_length = len(document)
|
||||
|
||||
for i in range(num_chunks):
|
||||
curr_pos = i * num_chunks
|
||||
to_pos = (
|
||||
curr_pos + num_chunks
|
||||
if curr_pos + num_chunks < document_length
|
||||
else document_length
|
||||
)
|
||||
text_chunk = document[curr_pos:to_pos]
|
||||
|
||||
collection.add(
|
||||
ids=[str(doc_uuid) + ":" + str(i)],
|
||||
documents=[text_chunk],
|
||||
)
|
||||
|
||||
return chunks
|
||||
|
||||
|
||||
# Setup
|
||||
|
||||
# Step 1: Get the text
|
||||
ppngx = PaperlessNGXService()
|
||||
docs = ppngx.get_data()
|
||||
texts = [doc["content"] for doc in docs]
|
||||
|
||||
# Step 2: Create chunks
|
||||
chunker = Chunker()
|
||||
|
||||
print(f"chunking {len(texts)} documents")
|
||||
for text in texts:
|
||||
chunker.chunk_document(document=text)
|
||||
|
||||
# Ask
|
||||
input = "How many teeth has Simba had removed?"
|
||||
response = ollama.embed(model="mxbai-embed-large", input=input)
|
||||
results = collection.query(query_texts=[input], n_results=1)
|
||||
print(results)
|
||||
# Generate
|
||||
output = ollama.generate(
|
||||
model="gemma3n:e4b",
|
||||
prompt=f"Using this data: {results}. Respond to this prompt: {input}",
|
||||
)
|
||||
|
||||
print(output["response"])
|
||||
7
pyproject.toml
Normal file
7
pyproject.toml
Normal file
@@ -0,0 +1,7 @@
|
||||
[project]
|
||||
name = "raggr"
|
||||
version = "0.1.0"
|
||||
description = "Add your description here"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.13"
|
||||
dependencies = []
|
||||
19
request.py
Normal file
19
request.py
Normal file
@@ -0,0 +1,19 @@
|
||||
import os
|
||||
import httpx
|
||||
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv()
|
||||
|
||||
|
||||
class PaperlessNGXService:
|
||||
def __init__(self):
|
||||
self.base_url = os.getenv("BASE_URL")
|
||||
self.token = os.getenv("PAPERLESS_TOKEN")
|
||||
self.url = f"http://{os.getenv("BASE_URL")}/api/documents/?query=simba"
|
||||
self.headers = {"Authorization": f"Token {os.getenv("PAPERLESS_TOKEN")}"}
|
||||
|
||||
def get_data(self):
|
||||
print(f"Getting data from: {self.url}")
|
||||
r = httpx.get(self.url, headers=self.headers)
|
||||
return r.json()["results"]
|
||||
Reference in New Issue
Block a user