| 572 | |
| 573 | @singleton |
| 574 | class SummaryEmbeddingHelper: |
| 575 | |
| 576 | openaiAccessRate = 0 |
| 577 | |
| 578 | def __init__(self, embedding_function: Optional[str] = None): |
| 579 | import chromadb |
| 580 | import streamlit as st |
| 581 | |
| 582 | self.client = chromadb.PersistentClient(path="./embedding") |
| 583 | |
| 584 | self.generated_ids = set() |
| 585 | |
| 586 | # Map{collectionId, collection} |
| 587 | self.collectionMap = {} |
| 588 | |
| 589 | # Map{collectionID, idIndex} |
| 590 | self.idIndexMap = {} |
| 591 | |
| 592 | self.api_key = st.session_state["openai_api_key"] |
| 593 | self.chatModel = ChatOpenAI(openai_api_key=self.api_key) |
| 594 | |
| 595 | if embedding_function is None: |
| 596 | self.embedding_model = embedding_functions.DefaultEmbeddingFunction() |
| 597 | elif embedding_function == "openai": |
| 598 | self.embedding_model = embedding_functions.OpenAIEmbeddingFunction(api_key=self.api_key) |
| 599 | |
| 600 | self.summaryPromptChinese = "请你对下面的段落进行总结,使用中文进行总结,具体的段落内容如下: {}" |
| 601 | self.summaryPromptEnglish = "Please summarize the following paragraph in English, the specific content of the " \ |
| 602 | "paragraph is as follows: {}" |
| 603 | |
| 604 | def from_documents(self, documents: list[Document], collectionId: Optional[str], embedding: Optional = OpenAIEmbeddings(), |
| 605 | persist_directory: Optional[str] = None) -> str: |
| 606 | if collectionId is None: |
| 607 | collectionId = self.generated_id() |
| 608 | collection = self.client.get_or_create_collection(collectionId) |
| 609 | idIndex = 0 |
| 610 | |
| 611 | texts = [] |
| 612 | ids = [] |
| 613 | embeddings = [] |
| 614 | for document in documents: |
| 615 | text = document.page_content |
| 616 | texts.append(text) |
| 617 | embeddings.append(self.summaryEmbedding(text)) |
| 618 | ids.append("id" + str(idIndex)) |
| 619 | idIndex += 1 |
| 620 | |
| 621 | collection.add( |
| 622 | documents=texts, |
| 623 | embeddings=embeddings, |
| 624 | # metadatas=[{"source": "notion"}, {"source": "google-docs"}], # filter on these! |
| 625 | ids=ids, |
| 626 | ) |
| 627 | |
| 628 | self.collectionMap[collectionId] = collection |
| 629 | self.idIndexMap[collectionId] = idIndex |
| 630 | return collectionId |
| 631 | |