feat OLLAMA_URL 切换到A6000

fix
This commit is contained in:
zhouchengrong
2024-11-08 14:35:23 +08:00
parent be7e121033
commit aa0db5006d

View File

@@ -6,6 +6,8 @@ from chromadb.config import Settings
from chromadb.utils.embedding_functions.ollama_embedding_function import OllamaEmbeddingFunction from chromadb.utils.embedding_functions.ollama_embedding_function import OllamaEmbeddingFunction
from tqdm import tqdm from tqdm import tqdm
from app.core.config import OLLAMA_URL
# 读取 csv 文件 # 读取 csv 文件
# csv_file_path = r'D:/Files/csv/output/output.csv' # csv_file_path = r'D:/Files/csv/output/output.csv'
# image_path = r'D:/images-clean' # image_path = r'D:/images-clean'
@@ -18,7 +20,7 @@ client = chromadb.Client(Settings(is_persistent=True, persist_directory="/vector
# client = chromadb.Client(Settings(is_persistent=True, persist_directory="D:/workspace/AiDLab/vector_db")) # client = chromadb.Client(Settings(is_persistent=True, persist_directory="D:/workspace/AiDLab/vector_db"))
# 创建集合 # 创建集合
# embedding_fn = OllamaEmbeddingFunction(url="http://localhost:11434/api/embeddings", model_name="mxbai-embed-large") # embedding_fn = OllamaEmbeddingFunction(url="http://localhost:11434/api/embeddings", model_name="mxbai-embed-large")
embedding_fn = OllamaEmbeddingFunction(url="http://10.1.1.240:11434/api/embeddings", model_name="mxbai-embed-large") embedding_fn = OllamaEmbeddingFunction(url=OLLAMA_URL, model_name="mxbai-embed-large")
# def create_collection(): # def create_collection():
@@ -67,7 +69,7 @@ embedding_fn = OllamaEmbeddingFunction(url="http://10.1.1.240:11434/api/embeddin
def query(gender, content): def query(gender, content):
collection = client.get_collection("sub_sketches_description", embedding_function=embedding_fn) collection = client.get_collection("sub_sketches_description", embedding_function=embedding_fn)
# 6. 查询相似内容 # 6. 查询相似内容
user_gender = gender # 用户输入的性别 user_gender = gender.lower() # 用户输入的性别
user_content = content # 用户输入的内容 user_content = content # 用户输入的内容
results = collection.query( results = collection.query(