From 33d1a05a77adbe6b28fbccf07f4f3b911bc4a228 Mon Sep 17 00:00:00 2001 From: Oyster Oil Date: Sun, 28 Apr 2024 18:36:53 +0800 Subject: [PATCH] Initial commit --- .github/workflows/build.yaml | 0 .gitignore | 5 ++ CITATION.cff | 11 ++++ LICENSE | 21 +++++++ README.md | 88 ++++++++++++++++++++++++++++ config.yaml | 7 +++ examples/api.py | 0 examples/simple.py | 59 +++++++++++++++++++ examples/with_config.py | 5 ++ pyproject.toml | 71 ++++++++++++++++++++++ ragbear/__init__.py | 4 ++ ragbear/__main__.py | 8 +++ ragbear/_entry.py | 45 ++++++++++++++ ragbear/answer.py | 5 ++ ragbear/api/__init__.py | 0 ragbear/api/app.py | 0 ragbear/config.py | 19 ++++++ ragbear/dababase/__init__.py | 0 ragbear/dababase/find.py | 42 +++++++++++++ ragbear/dababase/index/__init__.py | 0 ragbear/dababase/index/hnsw.py | 43 ++++++++++++++ ragbear/dababase/ingest.py | 13 ++++ ragbear/dababase/ingest/extract.py | 13 ++++ ragbear/dababase/schema.py | 13 ++++ ragbear/dababase/vendor/__init__.py | 0 ragbear/dababase/vendor/wiki_html.py | 16 +++++ ragbear/postprocess/rank.py | 3 + ragbear/query/prompts.py | 13 ++++ tests/test_extract.py | 19 ++++++ 29 files changed, 523 insertions(+) create mode 100644 .github/workflows/build.yaml create mode 100644 .gitignore create mode 100644 CITATION.cff create mode 100644 LICENSE create mode 100644 README.md create mode 100644 config.yaml create mode 100644 examples/api.py create mode 100644 examples/simple.py create mode 100644 examples/with_config.py create mode 100644 pyproject.toml create mode 100644 ragbear/__init__.py create mode 100644 ragbear/__main__.py create mode 100644 ragbear/_entry.py create mode 100644 ragbear/answer.py create mode 100644 ragbear/api/__init__.py create mode 100644 ragbear/api/app.py create mode 100644 ragbear/config.py create mode 100644 ragbear/dababase/__init__.py create mode 100644 ragbear/dababase/find.py create mode 100644 ragbear/dababase/index/__init__.py create mode 100644 ragbear/dababase/index/hnsw.py create mode 100644 ragbear/dababase/ingest.py create mode 100644 ragbear/dababase/ingest/extract.py create mode 100644 ragbear/dababase/schema.py create mode 100644 ragbear/dababase/vendor/__init__.py create mode 100644 ragbear/dababase/vendor/wiki_html.py create mode 100644 ragbear/postprocess/rank.py create mode 100644 ragbear/query/prompts.py create mode 100644 tests/test_extract.py diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml new file mode 100644 index 0000000..e69de29 diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..31703f9 --- /dev/null +++ b/.gitignore @@ -0,0 +1,5 @@ +dist/ +ragbear.egg-info/ +**/__pycache__ +data/ +.DS_Store diff --git a/CITATION.cff b/CITATION.cff new file mode 100644 index 0000000..b7f1b26 --- /dev/null +++ b/CITATION.cff @@ -0,0 +1,11 @@ +cff-version: 1.2.0 +message: "If you use this software, please cite it as below." +authors: +- family-names: "Lyonnet" + given-names: "Alex" +- family-names: "Yu" + given-names: "Shang" +title: "RAGBear - various RAG methods implemented in one package" +version: 0.0.1 +date-released: 2024-05-04 +url: "https://github.com/oysteroil/ragbear" diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..49c52a4 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2024 Oyster Oil + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 index 0000000..17ed506 --- /dev/null +++ b/README.md @@ -0,0 +1,88 @@ +# 🐻 Ragbear + +## Why + +RAG is simple, as simple as the following 10 lines of code: + +```py + +``` + +With `ragbear`, it's even simpler: + +```py +ragbear.query("Who is th father of Luke Skywalker?") +``` + +But a lot fancier, too: + +```py +ragbear.query( + "Who is the father of Luke Skywalker?", + model="http://localhost:8964", + source=[ + {"data": "http://localhost:8888", "type": "meilisearch"}, # query a meilisearch database + ], + rerank="newest", + rewrite="query2doc", + template=template, +) +``` + +But you will need to optimize towards your academic or commercial goals, you need a framework. +However, [existing solutions are too heavy](link to reddit). + +Enter Ragbear, we follow the above pattern closely, but give you options to swap each parts +of the pipeline. The code is concise and straightforward, no useless wrappers around wrappers. +Just read it. + +Unlike LangChain, we do not try to encapsulate every solution out there, because there is +actually not too much to wrap at all. We would like to call ragbear a patter, rather than +a framework. + +## Implemented algorithms + +With ragbear, we want to incorporate the latest research ideas and put it into production. +We have implemented the following algorithms, which you can easily apply in your app. + +- HyDE +- REPLUG +- Query2Doc + +For example, to use the REPLUG method, it's as easy as: + +```py +import ragbear + +ans = ragbear.query("Where was Steve Jobs born?", rewrite="replug") +``` + +## Data + +We support the following data query backend + +- DuckDuckGo, via python package `duckduckgo_search` +- Local Dataset, via ANN engines, aks vector databases + - hnswlib + sqlite + - faiss + sqlite + - pg_vector + +## Dependencies + +- An LLM API endpoint, you can use OpenAI API or local models with Ollama. +- An data backend, use `duckduckgo_search` or local dataset. + +## Citation + +If you find this project useful, please cite: + +``` +@misc{Charles2013, + author = {Alex Lyonnet, Shang Yu}, + title = {RAGBear - various RAG methods implemented in one package}, + year = {2024}, + publisher = {GitHub}, + journal = {GitHub repository}, + howpublished = {\url{https://github.com/oysteroil/ragbear}}, +} +``` diff --git a/config.yaml b/config.yaml new file mode 100644 index 0000000..e12a457 --- /dev/null +++ b/config.yaml @@ -0,0 +1,7 @@ +model: gpt-3 +embedding_model: gte +index_engine: hnswlib +query_rewrite: + - expand + - reform + diff --git a/examples/api.py b/examples/api.py new file mode 100644 index 0000000..e69de29 diff --git a/examples/simple.py b/examples/simple.py new file mode 100644 index 0000000..db8064c --- /dev/null +++ b/examples/simple.py @@ -0,0 +1,59 @@ +import ragbear + + +ragbear.query("Did Steve Jobs receive a PhD degree?") +# -> No. He did not receive a PhD degree. + + +# Use duckduckgo as the data source +prompt = "Did Steve Jobs receive a PhD degree?" +ragbear.query(prompt, source="duckduckgo") + + +# Use another model, you can use any openai compatible API endpoints +ragbear.query(prompt, model="http://localhost:8964") +ragbear.query(prompt, model="http://localhost:8964", openai_token="xxxxxx") + + +# Use a local datasource, you need to ingest the data first, see ingest.md for details +ragbear.query( + prompt, + source=[ + {"data": "/tmp/data/wiki.hnsw", "type": "hnsw", "max_results": 10}, # query local hnsw indexed files + {"data": "/tmp/data/wiki.bm25", "type": "bm25", "max_results": 10}, # query local bm25 indexed files + {"data": "http://localhost:8888", "type": "meilisearch", "max_results": 10}, # query a meilisearch database + ] +) + + +# Rerank the documents before generation +ragbear.query(prompt, rerank="similarity") +ragbear.query(prompt, rerank="oldest") +ragbear.query(prompt, rerank="newest") + + +# Rewrite the prompt before querying +ragbear.query(prompt, rewrite="query2doc") + + +# Use a different template for generation +template = """ +You are a helpful assistant, answer the following question with the references: + +Question: {question} +References: {refs} +""" +ragbear.query(prompt, template=template) + + +# Putting it all together +ragbear.query( + "Who is the father of Luke Skywalker?", + model="http://localhost:8964", + source=[ + {"data": "http://localhost:8888", "type": "meilisearch", "max_results": 10}, # query a meilisearch database + ], + rerank="newest", + rewrite="query2doc", + template=template, +) diff --git a/examples/with_config.py b/examples/with_config.py new file mode 100644 index 0000000..9572455 --- /dev/null +++ b/examples/with_config.py @@ -0,0 +1,5 @@ +import ragbear + +bear = ragbear.from_config("config.yaml") +ans = bear.query("Where was Steve Jobs born?", strategy="replug") +print(ans) diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..364a7cb --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,71 @@ +[project] +name = "ragbear" +version = "0.0.1a1" +authors = [{ name = "Oyster Oil", email = "infinitesheldon@gmail.com" }] +description = "Comprehensive RAG implementations" +license = { file = "LICENSE" } +dependencies = [ + "fastapi", + "torch", + "loguru", + "pyyaml", + "readability-lxml", + "markdownify", + "duckduckgo-search", +] +readme = "README.md" +requires-python = ">=3.8" +urls = { "repository" = "https://github.com/oysteroil/ragbear" } +classifiers = [ + "Development Status :: 4 - Beta", + "Intended Audience :: Developers", + "License :: OSI Approved :: MIT License", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", +] + + +[project.optional-dependencies] +dev = [ ] +build = [ ] +test = [ + "pytest>=8.1.1,<9.0", +] + + +[build-system] +requires = ["wheel", "setuptools"] +build-backend = "setuptools.build_meta" + + +[tool.setuptools] +packages = ["ragbear"] +package-data = {} + + +[tool.pytest.ini_options] +# pythonpath = [ "." ] +asyncio_mode = "auto" + + +[tool.ruff] +line-length = 88 + +[tool.ruff.lint] +select = [ + "E", # pycodestyle + "F", # Pyflakes + "UP", # pyupgrade + "B", # flake8-bugbear + "SIM", # flake8-simplify + "I", # isort +] + +[tool.mypy] +python_version = "3.8" +ignore_missing_imports = true +exclude = ["benchmark/", "docs/", "examples/", "scripts/", "tests/", "build/"] diff --git a/ragbear/__init__.py b/ragbear/__init__.py new file mode 100644 index 0000000..a0f93a5 --- /dev/null +++ b/ragbear/__init__.py @@ -0,0 +1,4 @@ +__all__ = ["query"] + + +from ._entry import query diff --git a/ragbear/__main__.py b/ragbear/__main__.py new file mode 100644 index 0000000..b8a4ae3 --- /dev/null +++ b/ragbear/__main__.py @@ -0,0 +1,8 @@ +import argparse + +parser = argparse.ArgumentParser() +subparser = parser.add_subparsers(dest="action") + +ingest = subparser.add_parser("ingest") +ingest.add_argument("--data") + diff --git a/ragbear/_entry.py b/ragbear/_entry.py new file mode 100644 index 0000000..e88219f --- /dev/null +++ b/ragbear/_entry.py @@ -0,0 +1,45 @@ +from contextlib import suppress +from typing import List, Literal, Optional, TypedDict, Union + +with suppress(ImportError): + pass + + +class SourceVendorType(TypedDict): + data: str + type: str + max_results: Optional[int] + + +SourceType = Union[Literal["duckduckgo"], List[SourceVendorType]] +RerankType = Literal["similarity", "oldest", "newest"] +RewriteType = Literal["query2doc"] + + +def query( + prompt: str, + model: Optional[str] = None, + *, + openai_token: Optional[str] = None, + source: Optional[SourceType] = None, + rerank: Optional[RerankType] = None, + rewrite: Optional[RewriteType] = None, + template: Optional[str] = None, + embedding_model: Optional[str] = None +): + """Query the LLM, and return the summarized result. + + Args: + prompt: the query string + model: the model to use, by default use local ollama server. + source: The data source for retrival + + """ + + # 1. preprocess the query + + # 2. find relevant docs + + # 3. query the LLM + + return diff --git a/ragbear/answer.py b/ragbear/answer.py new file mode 100644 index 0000000..b266540 --- /dev/null +++ b/ragbear/answer.py @@ -0,0 +1,5 @@ +from pydantic import BaseModel + + +class Answer(BaseModel): + text: str diff --git a/ragbear/api/__init__.py b/ragbear/api/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/ragbear/api/app.py b/ragbear/api/app.py new file mode 100644 index 0000000..e69de29 diff --git a/ragbear/config.py b/ragbear/config.py new file mode 100644 index 0000000..eeb92a3 --- /dev/null +++ b/ragbear/config.py @@ -0,0 +1,19 @@ +import yaml +from pathlib import Path +from typing import Literal, List + + +from pydantic import BaseModel + + +class Config(BaseModel): + index_engine: Literal["hnswlib", "faiss", "annoy"] + embedding_model: str + query_rewrite: List[Literal["expand"]] + + +def read_config(p: str | Path): + with open(p) as f: + return yaml.safe_load(f) + + diff --git a/ragbear/dababase/__init__.py b/ragbear/dababase/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/ragbear/dababase/find.py b/ragbear/dababase/find.py new file mode 100644 index 0000000..bf7abc6 --- /dev/null +++ b/ragbear/dababase/find.py @@ -0,0 +1,42 @@ +from typing import List, Optional +from pydantic import BaseModel, AnyUrl, Url +from sentence_transformers import SentenceTransformer +from duckduckgo_search import DDGS + +from ragbear.rank import noop_rank + + +model = SentenceTransformer("cyclone/simcse-chinese-roberta-wwm-ext") + +model.encode(["hehe", ""]) + + +class Doc(BaseModel): + title: str = "" + text: str = "" + full_text: str = "" + url: Optional[AnyUrl]= None + source: str = "" + source_type: str = "" + time_used: int + + +def find(query: str, recall="bm25"): + + # 1. vectorize the query string + vec = model.encode(str) + + # 2. find docs in the index + docs = index.search(vec) + + return docs + + +def find_bm25(query: str): + ... + + +def find_duckduckgo(query: str, max_results: int = 10) -> List[Doc]: + results = DDGS().text(query, max_results=max_results) + return docs + diff --git a/ragbear/dababase/index/__init__.py b/ragbear/dababase/index/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/ragbear/dababase/index/hnsw.py b/ragbear/dababase/index/hnsw.py new file mode 100644 index 0000000..8e91012 --- /dev/null +++ b/ragbear/dababase/index/hnsw.py @@ -0,0 +1,43 @@ +import hnswlib +import numpy as np + +DIM = 128 +num_elements = 10_000 + + +class HnswIndex: + + def __init__(self, dim: int, max_elements, num_threads: int=4, ef_construction: int = 200,): + """ + Args: + max_elements, the maximum number of elements (capacity). Will throw an exception if exceeded + + + """ + self.dim = dim + self.max_elements = max_elements + self.num_threads = num_threads + self.ef_construction = ef_construction + self._index = None + + def create(self): + index = hnswlib.Index(space="cosine", dim=self.dim) + index.init_index(max_elements=self.max_elements, ef_construction=200, M=16) + return index + + def load(self): + ... + + def save(self): + ... + + def add_items(self, items: list): + ... + + def query(self, q: str, k: int = 1): + labels, distance = self._index.knn_query(data=q, k=k) + labels = [[self.dict]] + ... + + + diff --git a/ragbear/dababase/ingest.py b/ragbear/dababase/ingest.py new file mode 100644 index 0000000..c2200c1 --- /dev/null +++ b/ragbear/dababase/ingest.py @@ -0,0 +1,13 @@ +import torch +from .index.hnsw import Index + + +def vectorize(s: str) -> torch.Tensor: + ... + + +def ingest(p: str, embedding_model: str = ""): + index = Index() + for doc in docs: + index.add(doc) + diff --git a/ragbear/dababase/ingest/extract.py b/ragbear/dababase/ingest/extract.py new file mode 100644 index 0000000..bcde257 --- /dev/null +++ b/ragbear/dababase/ingest/extract.py @@ -0,0 +1,13 @@ +from readability import Document +from markdownify import markdownify as md + + +def extract_md_from_html(s): + doc = Document(s) + html = doc.summary() + plain_text = md(html, default_title=True) + return plain_text + + +def chunked(s: str) -> list[str]: + return s.split("\n\n") diff --git a/ragbear/dababase/schema.py b/ragbear/dababase/schema.py new file mode 100644 index 0000000..7aed5e3 --- /dev/null +++ b/ragbear/dababase/schema.py @@ -0,0 +1,13 @@ +from pydantic import BaseModel + + +class Chunk(BaseModel): + id: int + text: str + article_id: int + + +class Article(BaseModel): + id: int + text: str + url: str diff --git a/ragbear/dababase/vendor/__init__.py b/ragbear/dababase/vendor/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/ragbear/dababase/vendor/wiki_html.py b/ragbear/dababase/vendor/wiki_html.py new file mode 100644 index 0000000..b419b3b --- /dev/null +++ b/ragbear/dababase/vendor/wiki_html.py @@ -0,0 +1,16 @@ +import zhconv +from pathlib import Path + +from ..index.hnsw import HnswIndex +from ..extract import extract_md_from_html + + +def import_all(p): + index = HnswIndex() + for file in Path().glob("**/*.html"): + with open(file) as f: + text = extract_md_from_html(f.read()) + text = text.replace("[编辑]", "") + simplfied = zhconv.convert(text) + index.add(simplfied) + diff --git a/ragbear/postprocess/rank.py b/ragbear/postprocess/rank.py new file mode 100644 index 0000000..f9511e0 --- /dev/null +++ b/ragbear/postprocess/rank.py @@ -0,0 +1,3 @@ +def noop_rank(docs): + return docs + diff --git a/ragbear/query/prompts.py b/ragbear/query/prompts.py new file mode 100644 index 0000000..fffa18b --- /dev/null +++ b/ragbear/query/prompts.py @@ -0,0 +1,13 @@ +from typing import List + +RAG_PROMPT = """请根据用户提问和参考资料进行回复，回复的内容控制在100字左右。 + +用户提问：{} + +参考材料： +{}""" + + + +def format(p: str, question: str, docs: List[str]): + return RAG_PROMPT.format(query, docs) diff --git a/tests/test_extract.py b/tests/test_extract.py new file mode 100644 index 0000000..ac802ab --- /dev/null +++ b/tests/test_extract.py @@ -0,0 +1,19 @@ +from ragbear.extract import extract_md_from_html + + +def test_extract_md_from_html(): + html = """ +

A New Era for Mixed Reality

+

Today we’re taking the next step toward our vision for a more open computing platform for the metaverse. We’re opening up the operating system powering our Meta Quest devices to third-party hardware makers, giving more choice to consumers and a larger ecosystem for developers to build for. We’re working with leading global technology companies to bring this new ecosystem to life and making it even easier for developers to build apps and reach their audiences on the platform.

+ """ + md = """A New Era for Mixed Reality +=========================== + + +Today we’re taking the next step toward our vision for a more open computing platform for the metaverse. We’re opening up the operating system powering our Meta Quest devices to third-party hardware makers, giving more choice to consumers and a larger ecosystem for developers to build for. We’re working with leading global technology companies to bring this new ecosystem to life and making it even easier for developers to build apps and reach their audiences on the platform. + + +""" + extracted = extract_md_from_html(html) + + assert extracted == md