103 lines
2.4 KiB
Python
103 lines
2.4 KiB
Python
import ollama
|
|
import os
|
|
from uuid import uuid4, UUID
|
|
|
|
from request import PaperlessNGXService
|
|
|
|
from math import ceil
|
|
|
|
import chromadb
|
|
|
|
from chromadb.utils.embedding_functions.ollama_embedding_function import (
|
|
OllamaEmbeddingFunction,
|
|
)
|
|
|
|
from dotenv import load_dotenv
|
|
|
|
client = chromadb.EphemeralClient()
|
|
collection = client.create_collection(name="docs")
|
|
|
|
load_dotenv()
|
|
|
|
|
|
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:
|
|
self.embedding_fx = OllamaEmbeddingFunction(
|
|
url=os.getenv("OLLAMA_URL", ""),
|
|
model_name="mxbai-embed-large",
|
|
)
|
|
|
|
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]
|
|
|
|
embedding = self.embedding_fx([text_chunk])
|
|
collection.add(
|
|
ids=[str(doc_uuid) + ":" + str(i)],
|
|
documents=[text_chunk],
|
|
embeddings=embedding,
|
|
)
|
|
|
|
return chunks
|
|
|
|
|
|
embedding_fx = OllamaEmbeddingFunction(
|
|
url=os.getenv("OLLAMA_URL", ""),
|
|
model_name="mxbai-embed-large",
|
|
)
|
|
|
|
# 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? Who is his current vet?"
|
|
embeddings = embedding_fx(input=[input])
|
|
results = collection.query(query_texts=[input], query_embeddings=embeddings)
|
|
print(results)
|
|
# Generate
|
|
output = ollama.generate(
|
|
model="gemma3n:e4b",
|
|
prompt=f"Using this data: {results}. Respond to this prompt: {input}",
|
|
)
|
|
|
|
print(output["response"])
|