github/DB-GPT
1
0
Fork 0
mirror of https://github.com/csunny/DB-GPT.git synced 2025-07-30 23:28:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

feature:source embedding

Browse Source
This commit is contained in:
chenketing 2023-05-11 16:06:18 +08:00
parent 365319a86c
commit 0c241dfaad
2 changed files with 24 additions and 85 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

43
pilot/source_embedding/pdf_embedding.py
View File

@ -1,16 +1,9 @@
#!/usr/bin/env python3 #!/usr/bin/env python3
# -*- coding: utf-8 -*- # -*- coding: utf-8 -*-
import json from typing import List
import os
from bs4 import BeautifulSoup from langchain.document_loaders import PyPDFLoader
from langchain.document_loaders import UnstructuredFileLoader, UnstructuredPDFLoader from langchain.schema import Document
from langchain.vectorstores import Milvus, Chroma
from pymilvus import connections
from pilot.server.vicuna_server import embeddings
from pilot.source_embedding.text_to_vector import TextToVector
# from vector_store import ESVectorStore
from pilot.source_embedding import SourceEmbedding, register from pilot.source_embedding import SourceEmbedding, register
@ -19,7 +12,7 @@ class PDFEmbedding(SourceEmbedding):
"""yuque embedding for read yuque document.""" """yuque embedding for read yuque document."""
def __init__(self, file_path, model_name, vector_store_config): def __init__(self, file_path, model_name, vector_store_config):
"""Initialize with YuqueLoader url.""" """Initialize with pdf path."""
self.file_path = file_path self.file_path = file_path
self.model_name = model_name self.model_name = model_name
self.vector_store_config = vector_store_config self.vector_store_config = vector_store_config
@ -27,28 +20,16 @@ class PDFEmbedding(SourceEmbedding):
@register @register
def read(self): def read(self):
"""Load from pdf path.""" """Load from pdf path."""
docs = [] loader = PyPDFLoader(self.file_path)
# loader = UnstructuredFileLoader(self.file_path) return loader.load()
loader = UnstructuredPDFLoader(self.file_path, mode="elements")
return loader.load()[0]
@register @register
def text_to_vector(self, docs): def data_process(self, documents: List[Document]):
"""Load from yuque url.""" i = 0
for doc in docs: for d in documents:
doc["vector"] = TextToVector.textToVector(doc["content"])[0] documents[i].page_content = d.page_content.replace(" ", "").replace("\n", "")
return docs i += 1
return documents
@register
def index_to_store(self, docs):
"""index into vector store."""
# vector_db = Milvus.add_texts(
# docs,
# embeddings,
# connection_args={"host": "127.0.0.1", "port": "19530"},
# )
db = Chroma.from_documents(docs, embeddings)
return Chroma.from_documents(docs, embeddings)

66
pilot/source_embedding/source_embedding.py
View File

@ -1,11 +1,10 @@
#!/usr/bin/env python3 #!/usr/bin/env python3
# -*- coding: utf-8 -*- # -*- coding: utf-8 -*-
import os
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from pymilvus import connections, FieldSchema, DataType, CollectionSchema from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import Chroma
from pilot.source_embedding.text_to_vector import TextToVector
from typing import List from typing import List
@ -30,9 +29,6 @@ class SourceEmbedding(ABC):
self.model_name = model_name self.model_name = model_name
self.vector_store_config = vector_store_config self.vector_store_config = vector_store_config
# Sub-classes should implement this method
# as return list(self.lazy_load()).
# This method returns a List which is materialized in memory.
@abstractmethod @abstractmethod
@register @register
def read(self) -> List[ABC]: def read(self) -> List[ABC]:
@ -49,61 +45,23 @@ class SourceEmbedding(ABC):
@register @register
def text_to_vector(self, docs): def text_to_vector(self, docs):
"""transform vector""" """transform vector"""
for doc in docs: pass
doc["vector"] = TextToVector.textToVector(doc["content"])[0]
return docs
@register @register
def index_to_store(self): def index_to_store(self, docs):
"""index to vector store""" """index to vector store"""
milvus = connections.connect( embeddings = HuggingFaceEmbeddings(model_name=self.model_name)
alias="default",
host='localhost',
port="19530"
)
doc_id = FieldSchema(
name="doc_id",
dtype=DataType.INT64,
is_primary=True,
)
doc_vector = FieldSchema(
name="doc_vector",
dtype=DataType.FLOAT_VECTOR,
dim=self.vector_store_config["dim"]
)
schema = CollectionSchema(
fields=[doc_id, doc_vector],
description=self.vector_store_config["description"]
)
@register persist_dir = os.path.join(self.vector_store_config["vector_store_path"],
def index_to_store(self): self.vector_store_config["vector_store_name"] + ".vectordb")
"""index to vector store""" vector_store = Chroma.from_documents(docs, embeddings, persist_directory=persist_dir)
milvus = connections.connect( vector_store.persist()
alias="default",
host='localhost',
port="19530"
)
doc_id = FieldSchema(
name="doc_id",
dtype=DataType.INT64,
is_primary=True,
)
doc_vector = FieldSchema(
name="doc_vector",
dtype=DataType.FLOAT_VECTOR,
dim=self.vector_store_config["dim"]
)
schema = CollectionSchema(
fields=[doc_id, doc_vector],
description=self.vector_store_config["description"]
)
def source_embedding(self): def source_embedding(self):
if 'read' in registered_methods: if 'read' in registered_methods:
text = self.read() text = self.read()
if 'process' in registered_methods: if 'data_process' in registered_methods:
self.process(text) text = self.data_process(text)
if 'text_split' in registered_methods: if 'text_split' in registered_methods:
self.text_split(text) self.text_split(text)
if 'text_to_vector' in registered_methods: if 'text_to_vector' in registered_methods: