bisheng-langchain 0.0.1__py3-none-any.whl

bisheng_langchain/document_loaders/parsers/test_image.py

@@ -0,0 +1,286 @@
+# flake8: noqa
+import io
+import json
+import logging
+import os
+import random
+import tempfile
+import time
+from abc import ABC
+from copy import deepcopy
+from pathlib import Path
+from typing import Any, Iterator, List, Mapping, Optional, Union
+from urllib.parse import urlparse
+
+import cv2
+import fitz
+import numpy as np
+import pypdfium2
+import requests
+from image import LayoutParser
+from langchain.document_loaders.blob_loaders import Blob
+
+
+def norm_rect(bbox):
+    x0 = np.min([bbox[0], bbox[2]])
+    x1 = np.max([bbox[0], bbox[2]])
+    y0 = np.min([bbox[1], bbox[3]])
+    y1 = np.max([bbox[1], bbox[3]])
+    return np.asarray([x0, y0, x1, y1])
+
+
+def merge_rects(bboxes):
+    x0 = np.min(bboxes[:, 0])
+    y0 = np.min(bboxes[:, 1])
+    x1 = np.max(bboxes[:, 2])
+    y1 = np.max(bboxes[:, 3])
+    return [x0, y0, x1, y1]
+
+
+def get_image_blobs(pages, pdf_reader, n, start=0):
+    blobs = []
+    for pg in range(start, start + n):
+        bytes_img = None
+        page = pages.load_page(pg)
+        mat = fitz.Matrix(1, 1)
+        try:
+            pm = page.get_pixmap(matrix=mat, alpha=False)
+            bytes_img = pm.getPNGData()
+        except Exception:
+            # some pdf input cannot get render image from fitz
+            page = pdf_reader.get_page(pg)
+            pil_image = page.render().to_pil()
+            img_byte_arr = io.BytesIO()
+            pil_image.save(img_byte_arr, format='PNG')
+            bytes_img = img_byte_arr.getvalue()
+
+            blobs.append(Blob(data=bytes_img))
+    return blobs
+
+
+def test():
+    file_path = './data/达梦数据库招股说明书_test_v1.pdf'
+    blob = Blob.from_path(file_path)
+    pages = None
+    image_blobs = []
+    with blob.as_bytes_io() as file_path:
+        pages = fitz.open(file_path)
+        pdf_reader = pypdfium2.PdfDocument(file_path, autoclose=True)
+        image_blobs = get_image_blobs(pages, pdf_reader)
+
+    assert len(image_blobs) == pages.page_count
+    layout = LayoutParser()
+    res = layout.parse(image_blobs[0])
+
+
+def draw_polygon(image, bbox, text=None, color=(255, 0, 0), thickness=1):
+    bbox = bbox.astype(np.int32)
+    is_rect = bbox.shape[0] == 4
+    if is_rect:
+        start_point = (bbox[0], bbox[1])
+        end_point = (bbox[2], bbox[3])
+        image = cv2.rectangle(image, start_point, end_point, color, thickness)
+    else:
+        polys = [bbox.astype(np.int32).reshape((-1, 1, 2))]
+        cv2.polylines(image, polys, True, color=color, thickness=thickness)
+        start_point = (polys[0][0, 0, 0], polys[0][0, 0, 1])
+
+    if text:
+        fontFace = cv2.FONT_HERSHEY_SIMPLEX
+        fontScale = 0.5
+        color = (0, 0, 255)
+        image = cv2.putText(image, text, start_point, fontFace, fontScale,
+                            color, 1)
+
+    return image
+
+
+def test_vis():
+    # file_path = './data/达梦数据库招股说明书_test_v1.pdf'
+    file_path = './data/pdf_input/《中国药典》2020年版 一部.pdf'
+    output_prefix = 'zhongguoyaodian_2020_v1'
+    start, end, n = 70, 80, 10
+    blob = Blob.from_path(file_path)
+    pages = None
+    image_blobs = []
+    with blob.as_bytes_io() as file_path:
+        pages = fitz.open(file_path)
+        pdf_reader = pypdfium2.PdfDocument(file_path, autoclose=True)
+        image_blobs = get_image_blobs(pages, pdf_reader, n, start)
+
+    assert len(image_blobs) == n
+
+    for i, blob in enumerate(image_blobs):
+        idx = i + start
+        # blob = image_blobs[2]
+        layout = LayoutParser()
+        out = layout.parse(blob)
+        res = json.loads(out[0].page_content)
+        bboxes = []
+        labels = []
+        for r in res:
+            bboxes.append(r['bbox'])
+            labels.append(str(r['category_id']))
+
+        bboxes = np.asarray(bboxes)
+
+        bytes_arr = np.frombuffer(blob.as_bytes(), dtype=np.uint8)
+        image = cv2.imdecode(bytes_arr, flags=1)
+        for bbox, text in zip(bboxes, labels):
+            image = draw_polygon(image, bbox, text)
+
+        outf = f'./data/{output_prefix}_layout_p{idx+1}_vis.png'
+        cv2.imwrite(outf, image)
+
+
+def order_by_tbyx(block_info, th=10):
+    """
+      block_info: [(b0, b1, b2, b3, text, x, y)+]
+      th: threshold of the position threshold
+    """
+    # sort using y1 first and then x1
+    res = sorted(block_info, key=lambda b: (b[1], b[0]))
+    for i in range(len(res) - 1):
+        for j in range(i, 0, -1):
+            # restore the order using the
+            if (abs(res[j + 1][1] - res[j][1]) < th
+                    and (res[j + 1][0] < res[j][0])):
+                tmp = deepcopy(res[j])
+                res[j] = deepcopy(res[j + 1])
+                res[j + 1] = deepcopy(tmp)
+            else:
+                break
+    return res
+
+
+def test_vis2():
+    # file_path = './data/达梦数据库招股说明书_test_v1.pdf'
+    file_path = './data/pdf_input/达梦数据库招股说明书.pdf'
+    output_prefix = 'dameng_pageblock'
+
+    start = 0
+    end = 10
+    n = end - start
+    blob = Blob.from_path(file_path)
+    pages = None
+    image_blobs = []
+    with blob.as_bytes_io() as file_path:
+        pages = fitz.open(file_path)
+        pdf_reader = pypdfium2.PdfDocument(file_path, autoclose=True)
+        image_blobs = get_image_blobs(pages, pdf_reader, n, start)
+
+    assert len(image_blobs) == pages.page_count
+
+    for i, blob in enumerate(image_blobs):
+        idx = i + start
+        page = pages.load_page(idx)
+
+        rect = page.rect
+        print('rect', rect)
+        o = 10
+        b0 = np.asarray([rect.x0 + o, rect.y0 + o, rect.x1 - o, rect.y1 - o])
+
+        bytes_arr = np.frombuffer(blob.as_bytes(), dtype=np.uint8)
+        image = cv2.imdecode(bytes_arr, flags=1)
+
+        image = draw_polygon(image, b0, '0.0')
+
+        textpage = page.get_textpage()
+        blocks = textpage.extractBLOCKS()
+        IMG_BLOCK_TYPE = 1
+
+        # blocks = order_by_tbyx(blocks)
+        bboxes = []
+        for off, b in enumerate(blocks):
+            label = 'text' if b[-1] != IMG_BLOCK_TYPE else 'image'
+            label = f'{label}-{off}'
+            print('block', b, label)
+            bbox = np.asarray([b[0], b[1], b[2], b[3]])
+            bboxes.append(bbox)
+
+            image = draw_polygon(image, bbox, label)
+
+        if bboxes:
+            b1 = merge_rects(np.asarray(bboxes))
+            b1 = np.asarray(b1)
+            image = draw_polygon(image, b1, '0.1')
+
+        outf = f'./data/{output_prefix}_p{idx}_vis.png'
+        cv2.imwrite(outf, image)
+
+
+def test_vis3():
+    file_path = './data/pdf_input/《中国药典》2020年版 一部.pdf'
+
+    start = 50
+    end = 60
+    n = end - start
+    output_prefix = 'zhongguoyaodian_2020_v1'
+
+    blob = Blob.from_path(file_path)
+    pages = None
+    image_blobs = []
+    with blob.as_bytes_io() as file_path:
+        pages = fitz.open(file_path)
+        pdf_reader = pypdfium2.PdfDocument(file_path, autoclose=True)
+        image_blobs = get_image_blobs(pages, pdf_reader, n, start=50)
+
+    assert len(image_blobs) == n
+
+    for i, blob in enumerate(image_blobs):
+        idx = i + start
+        page = pages.load_page(idx)
+
+        rect = page.rect
+        print('rect', rect)
+        o = 10
+        b0 = np.asarray([rect.x0 + o, rect.y0 + o, rect.x1 - o, rect.y1 - o])
+
+        bytes_arr = np.frombuffer(blob.as_bytes(), dtype=np.uint8)
+        image = cv2.imdecode(bytes_arr, flags=1)
+
+        image = draw_polygon(image, b0, '0.0')
+
+        rotation_matrix = np.asarray(page.rotation_matrix).reshape((3, 2))
+        c1 = (rotation_matrix[0, 0] - 1) <= 1e-6
+        c2 = (rotation_matrix[1, 1] - 1) <= 1e-6
+        is_rotated = c1 and c2
+
+        textpage = page.get_textpage()
+        blocks = textpage.extractBLOCKS()
+        IMG_BLOCK_TYPE = 1
+
+        # blocks = order_by_tbyx(blocks)
+        bboxes = []
+        for off, b in enumerate(blocks):
+            label = 'text' if b[-1] != IMG_BLOCK_TYPE else 'image'
+            label = f'{label}-{off}'
+            print('block', b, label)
+            bbox = np.asarray([b[0], b[1], b[2], b[3]])
+
+            aug_bbox = bbox.reshape((-1, 2))
+            padding = np.ones((len(aug_bbox), 1))
+            aug_bbox = np.hstack([aug_bbox, padding])
+            new_bbox = np.dot(aug_bbox, rotation_matrix).reshape(-1)
+
+            new_bbox = norm_rect(new_bbox)
+
+            print('new_bboxes', new_bbox)
+            bboxes.append(new_bbox)
+
+            image = draw_polygon(image, new_bbox, label)
+
+        print(bboxes)
+        if bboxes:
+            b1 = merge_rects(np.asarray(bboxes))
+            b1 = np.asarray(b1)
+            image = draw_polygon(image, b1, '0.1')
+
+        outf = f'./data/{output_prefix}_p{idx}_vis.png'
+        cv2.imwrite(outf, image)
+
+
+# test_vis3()
+# test_vis2()
+test_vis()
+# test()

bisheng_langchain/embeddings/__init__.py

@@ -0,0 +1,7 @@
+from .host_embedding import BGEZhEmbedding, GTEEmbedding, HostEmbeddings, ME5Embedding
+from .wenxin import WenxinEmbeddings
+
+__all__ = [
+    'WenxinEmbeddings', 'ME5Embedding', 'BGEZhEmbedding', 'GTEEmbedding',
+    'HostEmbeddings'
+]

bisheng_langchain/embeddings/host_embedding.py

@@ -0,0 +1,133 @@
+from __future__ import annotations
+
+import logging
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import requests
+from langchain.embeddings.base import Embeddings
+from langchain.utils import get_from_dict_or_env
+from pydantic import BaseModel, Extra, Field, root_validator
+from tenacity import (before_sleep_log, retry, retry_if_exception_type, stop_after_attempt,
+                      wait_exponential)
+
+logger = logging.getLogger(__name__)
+
+
+def _create_retry_decorator(
+        embeddings: HostEmbeddings) -> Callable[[Any], Any]:
+    min_seconds = 4
+    max_seconds = 10
+    # Wait 2^x * 1 second between each retry starting with
+    # 4 seconds, then up to 10 seconds, then 10 seconds afterwards
+    return retry(
+        reraise=True,
+        stop=stop_after_attempt(embeddings.max_retries),
+        wait=wait_exponential(multiplier=1, min=min_seconds, max=max_seconds),
+        retry=(retry_if_exception_type(Exception)),
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
+
+
+def embed_with_retry(embeddings: HostEmbeddings, **kwargs: Any) -> Any:
+    """Use tenacity to retry the embedding call."""
+    retry_decorator = _create_retry_decorator(embeddings)
+
+    @retry_decorator
+    def _embed_with_retry(**kwargs: Any) -> Any:
+        return embeddings.embed(**kwargs)
+
+    return _embed_with_retry(**kwargs)
+
+
+class HostEmbeddings(BaseModel, Embeddings):
+    """host embedding models.
+    """
+
+    client: Optional[Any]  #: :meta private:
+    """Model name to use."""
+    model: str = 'embedding-host'
+    host_base_url: str = None
+
+    deployment: Optional[str] = 'default'
+
+    embedding_ctx_length: Optional[int] = 6144
+    """The maximum number of tokens to embed at once."""
+    """Maximum number of texts to embed in each batch"""
+    max_retries: Optional[int] = 6
+    """Maximum number of retries to make when generating."""
+    request_timeout: Optional[Union[float, Tuple[float, float]]] = None
+    """Timeout in seconds for the OpenAPI request."""
+
+    model_kwargs: Optional[Dict[str, Any]] = Field(default_factory=dict)
+    """Holds any model parameters valid for `create` call not explicitly specified."""
+
+    verbose: Optional[bool] = False
+
+    class Config:
+        """Configuration for this pydantic object."""
+
+        extra = Extra.forbid
+
+    @root_validator()
+    def validate_environment(cls, values: Dict) -> Dict:
+        """Validate that api key and python package exists in environment."""
+        values['host_base_url'] = get_from_dict_or_env(values, 'host_base_url',
+                                                       'HostBaseUrl')
+
+        try:
+            values['client'] = requests.post
+        except AttributeError:
+            raise ValueError(
+                'Try upgrading it with `pip install --upgrade requests`.')
+        return values
+
+    @property
+    def _invocation_params(self) -> Dict:
+        api_args = {
+            'model': self.model,
+            'request_timeout': self.request_timeout,
+            **self.model_kwargs,
+        }
+        return api_args
+
+    def embed(self, texts: List[str], **kwargs) -> List[List[float]]:
+        emb_type = kwargs.get('type', 'raw')
+        inp = {'texts': texts, 'model': self.model, 'type': emb_type}
+        if self.verbose:
+            print('payload', inp)
+
+        url = f'{self.host_base_url}/{self.model}/infer'
+        outp = self.client(url=url, json=inp).json()
+        if outp['status_code'] != 200:
+            raise ValueError(
+                f"API returned an error: {outp['status_message']}")
+        return outp['embeddings']
+
+    def embed_documents(self,
+                        texts: List[str],
+                        chunk_size: Optional[int] = 0) -> List[List[float]]:
+        if not texts:
+            return []
+        """Embed search docs."""
+        texts = [text for text in texts if text]
+        embeddings = embed_with_retry(self, texts=texts, type='doc')
+        return embeddings
+
+    def embed_query(self, text: str) -> List[float]:
+        embeddings = embed_with_retry(self, texts=[text], type='query')
+        return embeddings[0]
+
+
+class ME5Embedding(HostEmbeddings):
+    model: str = 'multi-e5'
+    embedding_ctx_length: int = 512
+
+
+class BGEZhEmbedding(HostEmbeddings):
+    model: str = 'bge-zh'
+    embedding_ctx_length: int = 512
+
+
+class GTEEmbedding(HostEmbeddings):
+    model: str = 'gte'
+    embedding_ctx_length: int = 512

bisheng_langchain/embeddings/interface/__init__.py

@@ -0,0 +1,3 @@
+from .wenxin import EmbeddingClient as WenxinEmbeddingClient
+
+__all__ = ['WenxinEmbeddingClient']

bisheng_langchain/embeddings/interface/types.py

@@ -0,0 +1,23 @@
+from typing import Any, Dict, List, Union
+
+from pydantic import BaseModel
+
+
+class EmbeddingInput(BaseModel):
+    model: str
+    input: Union[str, List[str]]
+
+
+class Embedding(BaseModel):
+    object: str = 'embedding'
+    embedding: List[float]
+    index: int
+
+
+class EmbeddingOutput(BaseModel):
+    status_code: int
+    status_message: str = 'success'
+    object: str = None
+    data: List[Embedding] = []
+    model: str = None
+    usage: Dict[str, Any] = None

bisheng_langchain/embeddings/interface/wenxin.py

@@ -0,0 +1,86 @@
+import json
+
+import numpy as np
+import requests
+from requests.exceptions import HTTPError
+
+
+def get_access_token(api_key, sec_key):
+    url = (f'https://aip.baidubce.com/oauth/2.0/token?'
+           f'grant_type=client_credentials'
+           f'&client_id={api_key}&client_secret={sec_key}')
+
+    payload = json.dumps('')
+    headers = {
+        'Content-Type': 'application/json',
+        'Accept': 'application/json'
+    }
+
+    response = requests.request('POST', url, headers=headers, data=payload)
+    return response.json().get('access_token')
+
+
+class EmbeddingClient(object):
+
+    def __init__(self, api_key, sec_key, **kwargs):
+        self.api_key = api_key
+        self.sec_key = sec_key
+        self.ep_url = ('https://aip.baidubce.com/rpc/2.0/ai_custom/v1/'
+                       'wenxinworkshop/embeddings')
+        self.headers = {'Content-Type': 'application/json'}
+        self.max_text_tokens = 384
+        self.max_text_num = 16
+        self.drop_exceed_token = kwargs.get('drop_exceed_token', True)
+
+    def create(self, model, input, verbose=False, **kwargs):
+        texts = input
+        if isinstance(texts, str):
+            texts = [texts]
+
+        if self.drop_exceed_token:
+            texts = [t[:self.max_text_tokens] for t in texts]
+
+        cond = np.all([len(text) <= self.max_text_tokens for text in texts])
+        if not cond:
+            raise HTTPError('text exceed max token size 384')
+
+        token = get_access_token(self.api_key, self.sec_key)
+        endpoint = f'{self.ep_url}/{model}?access_token={token}'
+
+        def _call(sub_texts):
+            payload = json.dumps({'input': sub_texts})
+            response = requests.post(endpoint,
+                                     headers=self.headers,
+                                     data=payload)
+            status_message = 'success'
+            status_code = response.status_code
+            usage = {'prompt_tokens': 0, 'total_tokens': 0}
+            data = []
+            if status_code == 200:
+                try:
+                    info = json.loads(response.text)
+                    status_code = info.get('error_code', 200)
+                    status_message = info.get('error_msg', status_message)
+                    if status_code == 200:
+                        data = info['data']
+                        usage = info['usage']
+                    else:
+                        raise HTTPError(status_message)
+                except Exception as e:
+                    raise HTTPError(str(e))
+            else:
+                raise HTTPError('requests error')
+            return data, usage
+
+        data = []
+        usage = {'prompt_tokens': 0, 'total_tokens': 0}
+
+        for i in range(0, len(texts), self.max_text_num):
+            sub_texts = texts[i:(i + self.max_text_num)]
+            sub_data, sub_usage = _call(sub_texts)
+            data.extend(sub_data)
+            usage['prompt_tokens'] += sub_usage['prompt_tokens']
+            usage['total_tokens'] += sub_usage['total_tokens']
+
+        outp = dict(status_code=200, model=model, data=data, usage=usage)
+        return outp

bisheng_langchain/embeddings/wenxin.py

@@ -0,0 +1,139 @@
+from __future__ import annotations
+
+import logging
+# import warnings
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+# import numpy as np
+from langchain.embeddings.base import Embeddings
+from langchain.utils import get_from_dict_or_env
+from pydantic import BaseModel, Extra, Field, root_validator
+from requests.exceptions import HTTPError
+from tenacity import (before_sleep_log, retry, retry_if_exception_type, stop_after_attempt,
+                      wait_exponential)
+
+logger = logging.getLogger(__name__)
+
+
+def _create_retry_decorator(
+        embeddings: WenxinEmbeddings) -> Callable[[Any], Any]:
+    min_seconds = 4
+    max_seconds = 10
+    # Wait 2^x * 1 second between each retry starting with
+    # 4 seconds, then up to 10 seconds, then 10 seconds afterwards
+    return retry(
+        reraise=True,
+        stop=stop_after_attempt(embeddings.max_retries),
+        wait=wait_exponential(multiplier=1, min=min_seconds, max=max_seconds),
+        retry=(retry_if_exception_type(HTTPError)),
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
+
+
+def embed_with_retry(embeddings: WenxinEmbeddings, **kwargs: Any) -> Any:
+    """Use tenacity to retry the embedding call."""
+    retry_decorator = _create_retry_decorator(embeddings)
+
+    @retry_decorator
+    def _embed_with_retry(**kwargs: Any) -> Any:
+        return embeddings.embed(**kwargs)
+
+    return _embed_with_retry(**kwargs)
+
+
+class WenxinEmbeddings(BaseModel, Embeddings):
+    """Wenxin embedding models.
+
+    To use, the environment variable ``WENXIN_API_KEY`` and ``WENXIN_SECRET_KEY``
+    set with your API key or pass it as a named parameter to the constructor.
+
+    Example:
+        .. code-block:: python
+            from bisheng_langchain.embeddings import WenxinEmbeddings
+            wenxin_embeddings = WenxinEmbeddings(
+               wenxin_api_key="my-api-key",
+               wenxin_secret_key='xxx')
+
+    """
+
+    client: Optional[Any]  #: :meta private:
+    model: str = 'embedding-v1'
+
+    deployment: Optional[str] = 'default'
+    wenxin_api_key: Optional[str] = None
+    wenxin_secret_key: Optional[str] = None
+
+    embedding_ctx_length: Optional[int] = 6144
+    """The maximum number of tokens to embed at once."""
+    """Maximum number of texts to embed in each batch"""
+    max_retries: Optional[int] = 6
+    """Maximum number of retries to make when generating."""
+    request_timeout: Optional[Union[float, Tuple[float, float]]] = None
+    """Timeout in seconds for the OpenAPI request."""
+
+    model_kwargs: Optional[Dict[str, Any]] = Field(default_factory=dict)
+    """Holds any model parameters valid for `create` call not explicitly specified."""
+
+    class Config:
+        """Configuration for this pydantic object."""
+
+        extra = Extra.forbid
+
+    @root_validator()
+    def validate_environment(cls, values: Dict) -> Dict:
+        """Validate that api key and python package exists in environment."""
+        values['wenxin_api_key'] = get_from_dict_or_env(
+            values, 'wenxin_api_key', 'WENXIN_API_KEY')
+        values['wenxin_secret_key'] = get_from_dict_or_env(
+            values,
+            'wenxin_secret_key',
+            'WENXIN_SECRET_KEY',
+        )
+
+        api_key = values['wenxin_api_key']
+        sec_key = values['wenxin_secret_key']
+        try:
+            from .interface import WenxinEmbeddingClient
+            values['client'] = WenxinEmbeddingClient(api_key=api_key,
+                                                     sec_key=sec_key)
+        except AttributeError:
+            raise ValueError(
+                'Try upgrading it with `pip install --upgrade requests`.')
+        return values
+
+    @property
+    def _invocation_params(self) -> Dict:
+        wenxin_args = {
+            'model': self.model,
+            'request_timeout': self.request_timeout,
+            **self.model_kwargs,
+        }
+
+        return wenxin_args
+
+    def embed(self, texts: List[str]) -> List[List[float]]:
+        inp = {'input': texts, 'model': self.model}
+        outp = self.client.create(**inp)
+        if outp['status_code'] != 200:
+            raise ValueError(
+                f"Wenxin API returned an error: {outp['status_message']}")
+        return [e['embedding'] for e in outp['data']]
+
+    def embed_documents(self,
+                        texts: List[str],
+                        chunk_size: Optional[int] = 0) -> List[List[float]]:
+        embeddings = embed_with_retry(self, texts=texts)
+        return embeddings
+
+    def embed_query(self, text: str) -> List[float]:
+        """Call out to OpenAI's embedding endpoint for embedding query text.
+
+        Args:
+            text: The text to embed.
+
+        Returns:
+            Embedding for the text.
+        """
+
+        embeddings = embed_with_retry(self, texts=[text])
+        return embeddings[0]

bisheng_langchain/vectorstores/__init__.py

@@ -0,0 +1,3 @@
+from .elastic_keywords_search import ElasticKeywordsSearch
+
+__all__ = ['ElasticKeywordsSearch']