Adding getenv to main
This commit is contained in:
11
main.py
11
main.py
@@ -1,4 +1,5 @@
|
|||||||
import ollama
|
import ollama
|
||||||
|
import os
|
||||||
from uuid import uuid4, UUID
|
from uuid import uuid4, UUID
|
||||||
|
|
||||||
from request import PaperlessNGXService
|
from request import PaperlessNGXService
|
||||||
@@ -11,9 +12,13 @@ from chromadb.utils.embedding_functions.ollama_embedding_function import (
|
|||||||
OllamaEmbeddingFunction,
|
OllamaEmbeddingFunction,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
from dotenv import load_dotenv
|
||||||
|
|
||||||
client = chromadb.EphemeralClient()
|
client = chromadb.EphemeralClient()
|
||||||
collection = client.create_collection(name="docs")
|
collection = client.create_collection(name="docs")
|
||||||
|
|
||||||
|
load_dotenv()
|
||||||
|
|
||||||
|
|
||||||
class Chunk:
|
class Chunk:
|
||||||
def __init__(
|
def __init__(
|
||||||
@@ -34,7 +39,7 @@ class Chunk:
|
|||||||
class Chunker:
|
class Chunker:
|
||||||
def __init__(self) -> None:
|
def __init__(self) -> None:
|
||||||
self.embedding_fx = OllamaEmbeddingFunction(
|
self.embedding_fx = OllamaEmbeddingFunction(
|
||||||
url="http://localhost:11434",
|
url=os.getenv("OLLAMA_URL", ""),
|
||||||
model_name="mxbai-embed-large",
|
model_name="mxbai-embed-large",
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -67,7 +72,7 @@ class Chunker:
|
|||||||
|
|
||||||
|
|
||||||
embedding_fx = OllamaEmbeddingFunction(
|
embedding_fx = OllamaEmbeddingFunction(
|
||||||
url="http://localhost:11434",
|
url=os.getenv("OLLAMA_URL", ""),
|
||||||
model_name="mxbai-embed-large",
|
model_name="mxbai-embed-large",
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -84,7 +89,7 @@ for text in texts:
|
|||||||
chunker.chunk_document(document=text)
|
chunker.chunk_document(document=text)
|
||||||
|
|
||||||
# Ask
|
# Ask
|
||||||
input = "How many teeth has Simba had removed?"
|
input = "How many teeth has Simba had removed? Who is his current vet?"
|
||||||
embeddings = embedding_fx(input=[input])
|
embeddings = embedding_fx(input=[input])
|
||||||
results = collection.query(query_texts=[input], query_embeddings=embeddings)
|
results = collection.query(query_texts=[input], query_embeddings=embeddings)
|
||||||
print(results)
|
print(results)
|
||||||
|
|||||||
@@ -17,3 +17,8 @@ class PaperlessNGXService:
|
|||||||
print(f"Getting data from: {self.url}")
|
print(f"Getting data from: {self.url}")
|
||||||
r = httpx.get(self.url, headers=self.headers)
|
r = httpx.get(self.url, headers=self.headers)
|
||||||
return r.json()["results"]
|
return r.json()["results"]
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
pp = PaperlessNGXService()
|
||||||
|
print(pp.get_data()[0].keys())
|
||||||
|
|||||||
Reference in New Issue
Block a user