ryan/simbarag
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Only use OpenAI for embedding

Browse Source
This commit is contained in:
Ryan Chen
2025-10-14 20:06:32 -04:00
parent 376baccadb
commit b872750444
1 changed files with 5 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
chunker.py
View File

@@ -11,8 +11,6 @@ from dotenv import load_dotenv
from llm import LLMClient from llm import LLMClient
USE_OPENAI = os.getenv("OPENAI_API_KEY") != None
load_dotenv() load_dotenv()
ollama_client = Client(host=os.getenv("OLLAMA_HOST", "http://localhost:11434")) ollama_client = Client(host=os.getenv("OLLAMA_HOST", "http://localhost:11434"))
@@ -85,25 +83,16 @@ class Chunk:
class Chunker: class Chunker:
embedding_fx = OpenAIEmbeddingFunction(
api_key=os.getenv("OPENAI_API_KEY"),
model_name="text-embedding-3-small",
)
def __init__(self, collection) -> None: def __init__(self, collection) -> None:
self.collection = collection self.collection = collection
self.llm_client = LLMClient() self.llm_client = LLMClient()
def embedding_fx(self, inputs): def embedding_fx(self, inputs):
if self.llm_client.PROVIDER == "openai":
openai_embedding_fx = OpenAIEmbeddingFunction( openai_embedding_fx = OpenAIEmbeddingFunction(
api_key=os.getenv("OPENAI_API_KEY"), api_key=os.getenv("OPENAI_API_KEY"),
model_name="text-embedding-3-small", model_name="text-embedding-3-small",
) )
return openai_embedding_fx(inputs) return openai_embedding_fx(inputs)
else:
response = ollama_client.embed(model="mxbai-embed-large", input=inputs[0])
return response["embeddings"]
def chunk_document( def chunk_document(
self, self,