update codes

web_demo/backend/apis/api.py

+class API:
+    def call(self, **kwargs):
+        return self.__call__(**kwargs)
+
+    def __call__(self, **kwargs):
+        search_kwargs = {'query': kwargs['query'], 'freshness': 'month'}
+        search_res = self.search(**search_kwargs)
+        return search_res
+
+    @classmethod
+    def search(cls, **kwargs) -> list[str]:
+        raise NotImplementedError

web_demo/backend/apis/bing.py

+"""
+Bing Search
+"""
+import os
+
+import requests
+from backend.apis.api import API
+
+BING_API_KEY = os.getenv('BING_API_KEY')
+
+
+class BingSearchAPI(API):
+    def __init__(self):
+        self.url = "https://api.bing.microsoft.com/v7.0/search"
+
+    def search(self, query, freshness=None):
+        """
+        Search with bing
+
+        References: https://docs.microsoft.com/en-us/bing/search-apis/bing-web-search/overview
+        """
+        response = requests.get(
+            url=self.url,
+            headers={"Ocp-Apim-Subscription-Key": BING_API_KEY},
+            params={
+                "q": query,
+                "mkt": 'zh-CN',
+                "freshness": freshness,
+            },
+            timeout=10,
+        )
+        try:
+            json_content = response.json()
+            contexts = json_content['webPages']['value'][:4]
+            search_res = [{
+                "url": item['url'],
+                "title": item['name'],
+                "snippet": item['snippet']
+            } for item in contexts]
+            return search_res
+        except Exception as e:
+            print(f"Searching failed, caused by {e}")
+            return []

web_demo/backend/models/codegeex.py

+import torch
+from transformers import AutoTokenizer, AutoModel
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model_name_or_path = "THUDM/codegeex4-all-9b"
+tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=True)
+model = AutoModel.from_pretrained(
+    model_name_or_path,
+    trust_remote_code=True
+).to(device).eval()

web_demo/backend/services/chat.py

+import json
+
+from backend.apis.api import API
+from backend.apis.bing import BingSearchAPI
+from backend.models.codegeex import model, tokenizer
+from backend.utils.chat import build_model_input, SYS_PROMPT
+
+
+def chat(query: str, history: list[list[str]] = None):
+    if not history:
+        history = []
+
+    ans = ""
+
+    # Search with bing
+    api: API = BingSearchAPI()
+    search_res = api.call(query=query, history=history)
+    ans += "搜索结果".center(100, "-") + '\n'
+    ans += "```json\n" + json.dumps(search_res, indent=4, ensure_ascii=False) + "\n```\n"
+    yield ans
+
+    # Build model's input
+    inputs: str = build_model_input(query, search_res)
+
+    # Generate response
+    ans += "模型回复".center(100, "-") + '\n'
+    yield ans
+    response, _ = model.chat(
+        tokenizer,
+        query=inputs,
+        history=[{"role": "system", "content": SYS_PROMPT}],
+        max_new_tokens=1024,
+        temperature=0.2
+    )
+    yield ans + response

web_demo/backend/utils/chat.py

+SYS_PROMPT = """
+你将接收到一个用户提出的问题，并请撰写清晰、简洁且准确的答案。
+
+# Note
+- 您将获得与问题相关的多个上下文片段，每个上下文都以引用编号开头，例如[[citation:x]]，其中x是一个数字。如果适用，请使用上下文并在每个句子的末尾引用上下文。
+- 您的答案必须是正确的、准确的，并且以专家的身份使用无偏见和专业的语调来撰写。
+- 请你的回答限制在2千字以内，不要提供与问题无关的信息，也不要重复。
+- 请以引用编号的格式[[citation:x]]来引用上下文。如果一个句子来自多个上下文，请列出所有适用的引用，例如[[citation:3]][[citation:5]]。
+- 若所有上下文均不相关，请以自己的理解回答用户提出的问题，此时回答中可以不带引用编号。
+- 除了代码和特定的名称和引用外，您的答案必须使用与问题相同的语言来撰写。
+""".lstrip()
+
+
+def build_model_input(query, search_res):
+    """
+    Build model's input
+    :param query: user query
+    :param search_res: bing's search results
+    :return:
+    """
+    citations = "\n\n".join(
+        [f"[[citation:{i + 1}]]\n```markdown\n{item['snippet']}\n```" for i, item in enumerate(search_res)]
+    )
+    prompt = f"[引用]\n{citations}\n问：{query}\n"
+
+    return prompt

web_demo/main.py

+"""
+References: https://github.com/leptonai/search_with_lepton
+"""
+import gradio as gr
+
+from backend.services.chat import chat
+
+if __name__ == "__main__":
+    demo = gr.ChatInterface(chat).queue()
+    demo.launch(server_name="127.0.0.1", server_port=8080)

web_demo/requirements.txt

+accelerate==0.31.0
+gradio==4.26.0
+regex==2024.5.15
+requests==2.32.3
+tiktoken==0.7.0
+torch==2.3.1
+transformers==4.39.0