Spaces:
Running
Running
| from typing import List, Tuple | |
| import clickhouse_connect | |
| from sentence_transformers import SentenceTransformer | |
| from InstructorEmbedding import INSTRUCTOR | |
| emb_wiki = SentenceTransformer("sentence-transformers/paraphrase-multilingual-mpnet-base-v2") | |
| emb_arxiv = INSTRUCTOR('hkunlp/instructor-xl') | |
| class ArXivKnowledgeBase: | |
| def __init__(self, embedding: SentenceTransformer) -> None: | |
| self.db = clickhouse_connect.get_client( | |
| host='msc-4a9e710a.us-east-1.aws.staging.myscale.cloud', | |
| port=443, | |
| username='chatdata', | |
| password='myscale_rocks' | |
| ) | |
| self.embedding: SentenceTransformer = embedding | |
| self.table: str = 'default.ChatArXiv' | |
| self.embedding_col = "vector" | |
| self.must_have_cols: List[str] = ['id', 'abstract', 'authors', 'categories', 'comment', 'title', 'pubdate'] | |
| def __call__(self, subject: str, where_str: str = None, limit: int = 5) -> Tuple[str, int]: | |
| q_emb = self.embedding.encode(subject).tolist() | |
| q_emb_str = ",".join(map(str, q_emb)) | |
| if where_str: | |
| where_str = f"WHERE {where_str}" | |
| else: | |
| where_str = "" | |
| q_str = f""" | |
| SELECT dist, {','.join(self.must_have_cols)} | |
| FROM {self.table} | |
| {where_str} | |
| ORDER BY distance({self.embedding_col}, [{q_emb_str}]) | |
| AS dist ASC | |
| LIMIT {limit} | |
| """ | |
| docs = [r for r in self.db.query(q_str).named_results()] | |
| return '\n'.join([str(d) for d in docs]), len(docs) | |
| class WikiKnowledgeBase(ArXivKnowledgeBase): | |
| def __init__(self, embedding: SentenceTransformer) -> None: | |
| super().__init__(embedding) | |
| self.table: str = 'wiki.Wikipedia' | |
| self.embedding_col = "emb" | |
| self.must_have_cols: List[str] = ['text', 'title', 'views', 'url'] | |
| if __name__ == '__main__': | |
| # kb = ArXivKnowledgeBase(embedding=emb_arxiv) | |
| kb = WikiKnowledgeBase(embedding=emb_wiki) | |
| d = kb("When did Steven Jobs die?", "", 5) | |
| print(d) | |
| d = {"components": { | |
| "schemas": { | |
| "type": "object", | |
| "properties": { | |
| "todos":{ | |
| "type": "array", | |
| "items":{"type": "string"}, | |
| "description": "The list of todos.", | |
| } | |
| } | |
| } | |
| } | |
| } |