Create ERNIE-Bot-SDK.py

ERNIE-Bot-SDK.py

@@ -0,0 +1,737 @@
+#!/usr/bin/env python
+
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import math
+import os
+import time
+from collections.abc import Iterator
+from typing import List
+
+import faiss
+import gradio as gr
+import numpy as np
+import requests
+from tqdm import tqdm
+
+import erniebot as eb
+
+
+def parse_setup_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--port", type=int, default=8073)
+    args = parser.parse_args()
+    return args
+
+
+def create_ui_and_launch(args):
+    with gr.Blocks(title="ERNIE Bot SDK Demos", theme=gr.themes.Soft()) as blocks:
+        gr.Markdown("# ERNIE Bot SDK基础功能演示")
+        create_chat_completion_tab()
+        create_embedding_tab()
+        create_image_tab()
+        create_rag_tab()
+
+    blocks.launch(server_name="0.0.0.0", server_port=args.port)
+
+
+def create_chat_completion_tab():
+    def _infer(
+        ernie_model, content, state, top_p, temperature, api_type, access_key, secret_key, access_token
+    ):
+        access_key = access_key.strip()
+        secret_key = secret_key.strip()
+        access_token = access_token.strip()
+
+        if (access_key == "" or secret_key == "") and access_token == "":
+            raise gr.Error("需要填写正确的AK/SK或access token，不能为空")
+        if content.strip() == "":
+            raise gr.Error("输入不能为空，请在清空后重试")
+
+        auth_config = {
+            "api_type": api_type,
+        }
+        if access_key:
+            auth_config["ak"] = access_key
+        if secret_key:
+            auth_config["sk"] = secret_key
+        if access_token:
+            auth_config["access_token"] = access_token
+
+        content = content.strip().replace("<br>", "")
+        context = state.setdefault("context", [])
+        context.append({"role": "user", "content": content})
+        data = {
+            "messages": context,
+            "top_p": top_p,
+            "temperature": temperature,
+        }
+
+        if ernie_model == "chat_file":
+            response = eb.ChatFile.create(_config_=auth_config, **data, stream=False)
+        else:
+            response = eb.ChatCompletion.create(
+                _config_=auth_config, model=ernie_model, **data, stream=False
+            )
+
+        bot_response = response.result
+        context.append({"role": "assistant", "content": bot_response})
+        history = _get_history(context)
+        return None, history, context, state
+
+    def _regen_response(
+        ernie_model, state, top_p, temperature, api_type, access_key, secret_key, access_token
+    ):
+        """Regenerate response."""
+        context = state.setdefault("context", [])
+        if len(context) < 2:
+            raise gr.Error("请至少进行一轮对话")
+        context.pop()
+        user_message = context.pop()
+        return _infer(
+            ernie_model,
+            user_message["content"],
+            state,
+            top_p,
+            temperature,
+            api_type,
+            access_key,
+            secret_key,
+            access_token,
+        )
+
+    def _rollback(state):
+        """Roll back context."""
+        context = state.setdefault("context", [])
+        content = context[-2]["content"]
+        context = context[:-2]
+        state["context"] = context
+        history = _get_history(context)
+        return content, history, context, state
+
+    def _get_history(context):
+        history = []
+        for turn_idx in range(0, len(context), 2):
+            history.append([context[turn_idx]["content"], context[turn_idx + 1]["content"]])
+
+        return history
+
+    with gr.Tab("对话补全（Chat Completion）") as chat_completion_tab:
+        with gr.Row():
+            with gr.Column(scale=1):
+                api_type = gr.Dropdown(
+                    label="API Type", info="提供对话能力的后端平台", value="qianfan", choices=["qianfan", "aistudio"]
+                )
+                access_key = gr.Textbox(
+                    label="AK", info="用于访问后端平台的AK，如果设置了access token则无需设置此参数", type="password"
+                )
+                secret_key = gr.Textbox(
+                    label="SK", info="用于访问后端平台的SK，如果设置了access token则无需设置此参数", type="password"
+                )
+                access_token = gr.Textbox(
+                    label="Access Token", info="用于��问后端平台的access token，如果设置了AK、SK则无需设置此参数", type="password"
+                )
+                ernie_model = gr.Dropdown(
+                    label="Model", info="模型类型", value="ernie-bot", choices=["ernie-bot", "ernie-bot-turbo"]
+                )
+                top_p = gr.Slider(
+                    label="Top-p", info="控制采样范围，该参数越小生成结果越稳定", value=0.7, minimum=0, maximum=1, step=0.05
+                )
+                temperature = gr.Slider(
+                    label="Temperature",
+                    info="控制采样随机性，该参数越小生成结果越稳定",
+                    value=0.95,
+                    minimum=0.05,
+                    maximum=1,
+                    step=0.05,
+                )
+            with gr.Column(scale=4):
+                state = gr.State({})
+                context_chatbot = gr.Chatbot(label="对话历史")
+                input_text = gr.Textbox(label="消息内容", placeholder="请输入...")
+                with gr.Row():
+                    clear_btn = gr.Button("清空")
+                    rollback_btn = gr.Button("撤回")
+                    regen_btn = gr.Button("重新生成")
+                    send_btn = gr.Button("发送")
+                raw_context_json = gr.JSON(label="原始对话上下文信息")
+
+        api_type.change(
+            lambda api_type: {
+                "qianfan": (gr.update(visible=True), gr.update(visible=True)),
+                "aistudio": (gr.update(visible=False), gr.update(visible=False)),
+            }[api_type],
+            inputs=api_type,
+            outputs=[
+                access_key,
+                secret_key,
+            ],
+        )
+        chat_completion_tab.select(
+            lambda: (None, None, None, {}),
+            outputs=[
+                input_text,
+                context_chatbot,
+                raw_context_json,
+                state,
+            ],
+        )
+        input_text.submit(
+            _infer,
+            inputs=[
+                ernie_model,
+                input_text,
+                state,
+                top_p,
+                temperature,
+                api_type,
+                access_key,
+                secret_key,
+                access_token,
+            ],
+            outputs=[
+                input_text,
+                context_chatbot,
+                raw_context_json,
+                state,
+            ],
+        )
+        clear_btn.click(
+            lambda _: (None, None, None, {}),
+            inputs=clear_btn,
+            outputs=[
+                input_text,
+                context_chatbot,
+                raw_context_json,
+                state,
+            ],
+            show_progress=False,
+        )
+        rollback_btn.click(
+            _rollback,
+            inputs=[state],
+            outputs=[
+                input_text,
+                context_chatbot,
+                raw_context_json,
+                state,
+            ],
+            show_progress=False,
+        )
+        regen_btn.click(
+            _regen_response,
+            inputs=[
+                ernie_model,
+                state,
+                top_p,
+                temperature,
+                api_type,
+                access_key,
+                secret_key,
+                access_token,
+            ],
+            outputs=[
+                input_text,
+                context_chatbot,
+                raw_context_json,
+                state,
+            ],
+        )
+        send_btn.click(
+            _infer,
+            inputs=[
+                ernie_model,
+                input_text,
+                state,
+                top_p,
+                temperature,
+                api_type,
+                access_key,
+                secret_key,
+                access_token,
+            ],
+            outputs=[
+                input_text,
+                context_chatbot,
+                raw_context_json,
+                state,
+            ],
+        )
+
+
+def create_embedding_tab():
+    def _get_embeddings(text1, text2, api_type, access_key, secret_key, access_token):
+        access_key = access_key.strip()
+        secret_key = secret_key.strip()
+        access_token = access_token.strip()
+
+        if (access_key == "" or secret_key == "") and access_token == "":
+            raise gr.Error("需要填写正确的AK/SK或access token，不能为空")
+
+        auth_config = {
+            "api_type": api_type,
+        }
+        if access_key:
+            auth_config["ak"] = access_key
+        if secret_key:
+            auth_config["sk"] = secret_key
+        if access_token:
+            auth_config["access_token"] = access_token
+
+        if text1.strip() == "" or text2.strip() == "":
+            raise gr.Error("两个输入均不能为空")
+        embeddings = eb.Embedding.create(
+            _config_=auth_config,
+            model="ernie-text-embedding",
+            input=[text1.strip(), text2.strip()],
+        )
+        emb_0 = embeddings.rbody["data"][0]["embedding"]
+        emb_1 = embeddings.rbody["data"][1]["embedding"]
+        cos_sim = _calc_cosine_similarity(emb_0, emb_1)
+        cos_sim_text = f"## 两段文本余弦相似度: {cos_sim}"
+        return str(emb_0), str(emb_1), cos_sim_text
+
+    def _calc_cosine_similarity(vec_0, vec_1):
+        dot_result = float(np.dot(vec_0, vec_1))
+        denom = np.linalg.norm(vec_0) * np.linalg.norm(vec_1)
+        return 0.5 + 0.5 * (dot_result / denom) if denom != 0 else 0
+
+    with gr.Tab("语义向量（Embedding）"):
+        gr.Markdown("输入两段文本，分别获取两段文本的向量表示，并计算向量间的余弦相似度")
+        with gr.Row():
+            with gr.Column(scale=1):
+                api_type = gr.Dropdown(
+                    label="API Type", info="提供语义向量能力的后端平台", value="qianfan", choices=["qianfan", "aistudio"]
+                )
+                access_key = gr.Textbox(
+                    label="AK", info="用于访问后端平台的AK，如果设置了access token则无需设置此参数", type="password"
+                )
+                secret_key = gr.Textbox(
+                    label="SK", info="用于访问后端平台的SK，如果设置了access token则无需设置此参数", type="password"
+                )
+                access_token = gr.Textbox(
+                    label="Access Token", info="用于访问后端平台的access token，如果设置了AK、SK则无需设置此参数", type="password"
+                )
+            with gr.Column(scale=4):
+                with gr.Row():
+                    text1 = gr.Textbox(label="第一段文本", placeholder="输入第一段文本")
+                    text2 = gr.Textbox(label="第二段文本", placeholder="输入第二段文本")
+                cal_emb = gr.Button("生成向量")
+                cos_sim = gr.Markdown("## 余弦相似度: -")
+                with gr.Row():
+                    embedding1 = gr.Textbox(label="文本1向量结果")
+                    embedding2 = gr.Textbox(label="文本2向量结果")
+
+        api_type.change(
+            lambda api_type: {
+                "qianfan": (gr.update(visible=True), gr.update(visible=True)),
+                "aistudio": (gr.update(visible=False), gr.update(visible=False)),
+            }[api_type],
+            inputs=api_type,
+            outputs=[
+                access_key,
+                secret_key,
+            ],
+        )
+        cal_emb.click(
+            _get_embeddings,
+            inputs=[
+                text1,
+                text2,
+                api_type,
+                access_key,
+                secret_key,
+                access_token,
+            ],
+            outputs=[
+                embedding1,
+                embedding2,
+                cos_sim,
+            ],
+        )
+
+
+def create_image_tab():
+    def _gen_image(prompt, w_and_h, api_type, access_key, secret_key, access_token):
+        access_key = access_key.strip()
+        secret_key = secret_key.strip()
+        access_token = access_token.strip()
+
+        if (access_key == "" or secret_key == "") and access_token == "":
+            raise gr.Error("需要填写正确的AK/SK或access token，不能为空")
+        if prompt.strip() == "":
+            raise gr.Error("输入不能为空")
+
+        auth_config = {
+            "api_type": api_type,
+        }
+        if access_key:
+            auth_config["ak"] = access_key
+        if secret_key:
+            auth_config["sk"] = secret_key
+        if access_token:
+            auth_config["access_token"] = access_token
+
+        timestamp = int(time.time())
+        w, h = [int(x) for x in w_and_h.strip().split("x")]
+
+        response = eb.Image.create(
+            _config_=auth_config,
+            model="ernie-vilg-v2",
+            prompt=prompt,
+            width=w,
+            height=h,
+            version="v2",
+            image_num=1,
+        )
+        img_url = response.data["sub_task_result_list"][0]["final_image_list"][0]["img_url"]
+        res = requests.get(img_url)
+        with open(f"{timestamp}.jpg", "wb") as f:
+            f.write(res.content)
+        return f"{timestamp}.jpg"
+
+    with gr.Tab("文生图（Image Generation）"):
+        with gr.Row():
+            with gr.Column(scale=1):
+                api_type = gr.Dropdown(
+                    label="API Type", info="提供文生图能力的后端平台", value="yinian", choices=["yinian"]
+                )
+                access_key = gr.Textbox(
+                    label="AK", info="用于访问后端平台的AK，如果设置了access token则无需设置此参数", type="password"
+                )
+                secret_key = gr.Textbox(
+                    label="SK", info="用于访问后端平台的SK，如果设置了access token则无需设置此参数", type="password"
+                )
+                access_token = gr.Textbox(
+                    label="Access Token", info="用于访问后端平台的access token，如果设置了AK、SK则无需设置此参数", type="password"
+                )
+            with gr.Column(scale=4):
+                with gr.Row():
+                    prompt = gr.Textbox(label="Prompt", placeholder="输入用于生成图片的prompt，例如： 生成一朵玫瑰花")
+                    w_and_h = gr.Dropdown(
+                        label="分辨率",
+                        value="512x512",
+                        choices=[
+                            "512x512",
+                            "640x360",
+                            "360x640",
+                            "1024x1024",
+                            "1280x720",
+                            "720x1280",
+                            "2048x2048",
+                            "2560x1440",
+                            "1440x2560",
+                        ],
+                    )
+                submit_btn = gr.Button("生成图片")
+                image_show_zone = gr.Image(label="图片生成结果", type="filepath", show_download_button=True)
+
+        submit_btn.click(
+            _gen_image,
+            inputs=[
+                prompt,
+                w_and_h,
+                api_type,
+                access_key,
+                secret_key,
+                access_token,
+            ],
+            outputs=image_show_zone,
+        )
+
+
+def create_rag_tab():
+    REF_HTML = """
+
+    <details style="border: 1px solid #ccc; padding: 10px; border-radius: 4px; margin-bottom: 4px">
+        <summary style="display: flex; align-items: center; font-weight: bold;">
+            <span style="margin-right: 10px;">[{index}] {title}</span>
+            <a style="text-decoration: none; background: none !important;" target="_blank">
+                <!--[Here should be a link icon]-->
+                <i style="border: solid #000; border-width: 0 2px 2px 0; display: inline-block; padding: 3px;
+                transform:rotate(-45deg);-webkit-transform(-45deg)">
+                </i>
+            </a>
+        </summary>
+        <p style="margin-top: 10px;">{text}</p>
+    </details>
+
+    """
+
+    PROMPT_TEMPLATE = """基于以下已知信息，请简洁并专业地回答用户的问题。
+如果无法从中得到答案，请说 '根据已知信息无法回答该问题' 或 '没有提供足够的相关信息'。不允许在答案中添加编造成分。
+你可以参考以下文章:
+{DOCS}
+问题：{QUERY}
+回答："""
+
+    _CONFIG = {
+        "ernie_model": "",
+        "api_type": "",
+        "AK": "",
+        "SK": "",
+        "access_token": "",
+        "top_p": 0.7,
+        "temperature": 0.95,
+    }
+
+    def split_by_len(texts: List[str], split_token: int = 384) -> List[str]:
+        """
+        Split the knowledge base docs into chunks by length.
+
+        Args:
+            texts (List[str]): Knowledge Base Texts.
+            split_token (int, optional): The max length supported by ernie-text-embedding. Default to 384.
+
+        Returns:
+            List[str]: Doc Chunks.
+        """
+        chunk = []
+        for text in texts:
+            idx = 0
+            while idx + split_token < len(text):
+                temp_text = text[idx : idx + split_token]
+                next_idx = temp_text.rfind("。") + 1
+                if next_idx != 0:  # If this slice doesn't have a period, add the whole sentence.
+                    chunk.append(temp_text[:next_idx])
+                    idx = idx + next_idx
+                else:
+                    chunk.append(temp_text)
+                    idx = idx + split_token
+
+            chunk.append(text[idx:])
+        return chunk
+
+    def _get_embedding_doc(word: List[str]) -> List[float]:
+        """
+        Get the embedding of a list of words.
+
+        Args:
+            word (List[str]): Words to get embedding.
+
+        Returns:
+            List[float]: Embedding List of the words.
+        """
+        if (_CONFIG["AK"] == "" or _CONFIG["SK"] == "") and _CONFIG["access_token"] == "":
+            raise gr.Error("需要填写正确的AK/SK或access token，不能为空")
+
+        embedding: List[float]
+        if len(word) <= 16:
+            resp = eb.Embedding.create(model="ernie-text-embedding", input=word)
+            assert not isinstance(resp, Iterator)
+            embedding = resp.get_result()
+        else:
+            size = len(word)
+            embedding = []
+            for i in tqdm(range(math.ceil(size / 16))):
+                temp_result = eb.Embedding.create(
+                    model="ernie-text-embedding", input=word[i * 16 : (i + 1) * 16]
+                )
+                assert not isinstance(temp_result, Iterator)
+                embedding.extend(temp_result.get_result())
+                time.sleep(1)
+        return embedding
+
+    def l2_normalization(embedding: np.ndarray) -> np.ndarray:
+        "Vector Normalization by l2 norm"
+        if embedding.ndim == 1:
+            return embedding / np.linalg.norm(embedding).reshape(-1, 1)
+        else:
+            return embedding / np.linalg.norm(embedding, axis=1).reshape(-1, 1)
+
+    def find_related_doc(
+        query: str, origin_chunk: List[str], index_ip: faiss.swigfaiss.IndexFlatIP, top_k: int = 5
+    ) -> tuple[str, List[int]]:
+        """
+        Fin top_k similar documents.
+
+        Args:
+            query (str): user query.
+            origin_chunk (List[str]): Knowledge Base Doc.
+            index_ip (faiss.swigfaiss.IndexFlatIP): Vector DB index。
+            top_k (int, optional): Return top_k most similar documents. Default to 5.
+
+        Returns:
+            str, List[int]: The most similar documents and their index.
+        """
+
+        D, Idx = index_ip.search(np.array(_get_embedding_doc([query])), top_k)
+        top_k_similar = Idx.tolist()[0]
+
+        res = ""
+        ref_lis = []
+        for i in range(top_k):
+            res += f"[参考文章{i+1}]:{origin_chunk[top_k_similar[i]]}" + "\n\n"
+            ref_lis.append(origin_chunk[top_k_similar[i]])
+        return res, ref_lis
+
+    def process_uploaded_file(files: List[str], *args: object) -> str:
+        """
+        Args:
+            files: Files path
+            _CONFIG: Config
+        """
+        _update_config(*args)
+
+        content = []
+        for file in files:
+            with open(file, "r") as f:
+                content.append(f.read())
+
+        doc_chunk = split_by_len(content)
+
+        doc_embedding = _get_embedding_doc(doc_chunk)
+        assert len(doc_embedding) == len(doc_chunk), "shape mismatch"
+        doc_embedding_arr = l2_normalization(np.array(doc_embedding))
+
+        index_ip = faiss.IndexFlatIP(doc_embedding_arr.shape[1])
+        index_ip.add(doc_embedding_arr)
+
+        temp_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data")
+        if not os.path.exists(temp_path):
+            os.makedirs(temp_path)
+
+        faiss.write_index(index_ip, os.path.join(temp_path, "knowledge_embedding.index"))
+        with open(os.path.join(temp_path, "knowledge.txt"), "w") as f:
+            for chunk in doc_chunk:
+                f.write(repr(chunk) + "\n")
+
+        return "已完成向量知识库搭建"
+
+    def get_ans(query: str, *args: object) -> tuple[str, str]:
+        _update_config(*args)
+
+        if (_CONFIG["AK"] == "" or _CONFIG["SK"] == "") and _CONFIG["access_token"] == "":
+            raise gr.Error("需要填写正确的AK/SK或access token，不能为空")
+        temp_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data")
+        doc_chunk = []
+        with open(os.path.join(temp_path, "knowledge.txt"), "r") as f:
+            for line in f:
+                doc_chunk.append(eval(line))
+        index_ip = faiss.read_index(os.path.join(temp_path, "knowledge_embedding.index"))
+        related_doc, references = find_related_doc(query, doc_chunk, index_ip)
+
+        refs = []
+        for i in range(len(references)):
+            temp_dict = {
+                "title": f"Reference{i+1}",
+                "text": references[i],
+            }
+            refs.append(temp_dict)
+
+        resp = eb.ChatCompletion.create(
+            model=_CONFIG["ernie_model"],
+            messages=[{"role": "user", "content": PROMPT_TEMPLATE.format(DOCS=related_doc, QUERY=query)}],
+            top_p=_CONFIG["top_p"],
+            temperature=_CONFIG["temperature"],
+        )
+        assert not isinstance(resp, Iterator)
+        answer = resp.get_result()
+
+        return answer, "<h3>References (Click to Expand)</h3>" + "\n".join(
+            [REF_HTML.format(**item, index=idx + 1) for idx, item in enumerate(refs)]
+        )
+
+    def _update_config(*args: object):
+        eb.api_type = args[1]
+        eb.access_token = args[2]
+        eb.AK = args[3]
+        eb.SK = args[4]
+
+        _CONFIG.update(
+            {
+                "ernie_model": args[0],
+                "api_type": args[1],
+                "access_token": args[2],
+                "AK": args[3],
+                "SK": args[4],
+                "top_p": args[5],
+                "temperature": args[6],
+            }
+        )
+        # print(_CONFIG)
+
+    with gr.Tab("知识库问答(Retrieval Augmented QA)"):
+        # gr.Markdown("# 文心大模型RAG问答DEMO")
+        with gr.Tabs():
+            with gr.TabItem("设置栏"):
+                with gr.Row():
+                    with gr.Column():
+                        file_upload = gr.Files(file_types=["txt"], label="目前仅支持txt格式文件")
+                        chat_box = gr.Textbox(show_label=False)
+                    with gr.Column():
+                        ernie_model = gr.Dropdown(
+                            label="Model",
+                            info="模型类型",
+                            value="ernie-bot-4",
+                            choices=["ernie-bot-4", "ernie-bot-turbo", "ernie-bot"],
+                        )
+                        api_type = gr.Dropdown(
+                            label="API Type",
+                            info="提供���话能力的后端平台",
+                            value="aistudio",
+                            choices=["aistudio", "qianfan"],
+                        )
+                        access_token = gr.Textbox(
+                            label="Access Token",
+                            info="用于访问后端平台的access token，如果选择aistudio，则需设置此参数",
+                            type="password",
+                        )
+                        access_key = gr.Textbox(
+                            label="AK", info="用于访问千帆平台的AK，如果选择qianfan，则需设置此参数", type="password"
+                        )
+                        secret_key = gr.Textbox(
+                            label="SK", info="用于访问千帆平台的SK，如果选择qianfan，则需设置此参数", type="password"
+                        )
+                        top_p = gr.Slider(
+                            label="Top-p",
+                            info="控制采样范围，该参数越小生成结果越稳定",
+                            value=0.7,
+                            step=0.05,
+                            minimum=0,
+                            maximum=1,
+                        )
+                        temperature = gr.Slider(
+                            label="temperature",
+                            info="控制采样随机性，该参数越小生成结果越稳定",
+                            value=0.95,
+                            step=0.05,
+                            maximum=1,
+                            minimum=0,
+                        )
+
+            with gr.TabItem("问答栏"):
+                with gr.Row():
+                    query_box = gr.Textbox(show_label=False, placeholder="Enter question and press ENTER")
+
+                answer_box = gr.Textbox(show_label=False, value="", lines=5)
+                ref_boxes = gr.HTML(label="References")
+
+        query_box.submit(
+            get_ans,
+            [query_box, ernie_model, api_type, access_token, access_key, secret_key, top_p, temperature],
+            [answer_box, ref_boxes],
+        )
+        file_upload.upload(
+            process_uploaded_file,
+            [file_upload, ernie_model, api_type, access_token, access_key, secret_key, top_p, temperature],
+            chat_box,
+        )
+
+
+if __name__ == "__main__":
+    args = parse_setup_args()
+    create_ui_and_launch(args)