Upload 13 files

ChatHaruhi/BaseDB.py

@@ -0,0 +1,27 @@
+# BaseDB.py
+
+from abc import ABC, abstractmethod
+
+class BaseDB(ABC):
+
+    @abstractmethod
+    def init_db(self):
+        pass
+    
+    @abstractmethod
+    def save(self, file_path):
+        pass
+
+    @abstractmethod
+    def load(self, file_path):
+        pass
+
+    @abstractmethod
+    def search(self, vector, n_results):
+        pass
+
+    @abstractmethod
+    def init_from_docs(self, vectors, documents):
+        pass
+
+    

ChatHaruhi/BaseLLM.py

@@ -0,0 +1,56 @@
+# ChatHaruhi: Reviving Anime Character in Reality via Large Language Model
+#
+# ChatHaruhi 2.0, built by Cheng Li and Weishi Mi
+#
+# [email protected], [email protected]
+# 
+# Weishi Mi is a second-year graduate student at Tsinghua University, majoring in computer science.
+# Weishi Mi is pursuing a job or a PhD position, which who will be available next year
+# 
+# homepage https://github.com/LC1332/Chat-Haruhi-Suzumiya
+# 
+# ChatHaruhi is a chatbot that can revive anime characters in reality.
+# the 2.0 version was built by Cheng Li and Weishi Mi.
+# 
+# Please cite our paper if you use this code for research: 
+#
+# @misc{li2023chatharuhi,
+#       title={ChatHaruhi: Reviving Anime Character in Reality via Large Language Model}, 
+#       author={Cheng Li and Ziang Leng and Chenxi Yan and Junyi Shen and Hao Wang and Weishi MI and Yaying Fei and Xiaoyang Feng and Song Yan and HaoSheng Wang and Linkang Zhan and Yaokai Jia and Pingyu Wu and Haozhen Sun},
+#       year={2023},
+#       eprint={2308.09597},
+#       archivePrefix={arXiv},
+#       primaryClass={cs.CL}
+# }
+from abc import ABC, abstractmethod
+
+class BaseLLM(ABC):
+
+    def __init__(self):
+        pass
+    
+    @abstractmethod
+    def initialize_message(self):
+        pass
+
+    @abstractmethod    
+    def ai_message(self, payload):
+        pass
+
+    @abstractmethod
+    def system_message(self, payload):
+        pass
+
+    @abstractmethod
+    def user_message(self, payload):
+        pass
+
+    @abstractmethod
+    def get_response(self):
+        pass
+
+    @abstractmethod
+    def print_prompt(self):
+        pass
+
+

ChatHaruhi/ChatGLM2GPT.py

@@ -0,0 +1,69 @@
+import os
+from transformers import AutoTokenizer, AutoModel
+from peft import LoraConfig, get_peft_model
+from peft import PeftModel, PeftConfig
+from .BaseLLM import BaseLLM
+import torch 
+
+tokenizer_GLM = None
+model_GLM = None
+
+def initialize_GLM2LORA():
+    pass
+    global tokenizer_GLM
+    global model_GLM
+
+    if tokenizer_GLM == None and model_GLM == None:
+        tokenizer_GLM = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
+        model_GLM = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True).half().cuda()
+
+        config = LoraConfig(
+            r=16,
+            lora_alpha=32,
+            inference_mode=True,
+            lora_dropout=0.05,
+            #bias="none",
+            task_type="CAUSAL_LM"
+        )
+
+        model_GLM = PeftModel.from_pretrained(model_GLM, "silk-road/Chat-Haruhi-Fusion_B")
+    return model_GLM, tokenizer_GLM
+
+def GLM_tokenizer(text):
+    return len(tokenizer_GLM.encode(text))
+
+class ChatGLM2GPT(BaseLLM):
+    def __init__(self, model = "haruhi-fusion"):
+        super(ChatGLM2GPT, self).__init__()
+        if model == "glm2-6b":
+            self.tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
+            self.model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True).half().cuda()
+        if model == "haruhi-fusion":
+            self.model, self.tokenizer = initialize_GLM2LORA()
+        else:
+            raise Exception("Unknown GLM model")
+        self.messages = ""
+
+    def initialize_message(self):
+        self.message = ""
+
+    def ai_message(self, payload):
+        self.messages = self.messages + "\n " + payload 
+
+    def system_message(self, payload):
+        self.messages = self.messages + "\n " + payload 
+
+    def user_message(self, payload):
+        self.messages = self.messages + "\n " + payload 
+
+    def get_response(self):
+        with torch.no_grad():
+            response, history = self.model.chat(self.tokenizer, self.messages, history=[])
+            # print(response)
+        return response
+        
+    def print_prompt(self):
+        print(type(self.messages))
+        print(self.messages)
+
+    

ChatHaruhi/ChatHaruhi.py

@@ -0,0 +1,272 @@
+from .ChromaDB import ChromaDB
+import os
+
+from .utils import luotuo_openai_embedding, tiktokenizer
+
+from .utils import response_postprocess
+
+class ChatHaruhi:
+
+    def __init__(self, system_prompt = None, \
+                 role_name = None, \
+                 story_db=None, story_text_folder = None, \
+                 llm = 'openai', \
+                 embedding = 'luotuo_openai', \
+                 max_len_story = None, max_len_history = None,
+                 verbose = False):
+        super(ChatHaruhi, self).__init__()
+        self.verbose = verbose
+
+        # constants
+        self.story_prefix_prompt = "Classic scenes for the role are as follows:\n"
+        self.k_search = 19
+        self.narrator = ['旁白', '', 'scene','Scene','narrator' , 'Narrator']
+        self.dialogue_divide_token = '\n###\n'
+        self.dialogue_bra_token = '「'
+        self.dialogue_ket_token = '」'
+
+        if system_prompt:
+            self.system_prompt = self.check_system_prompt( system_prompt )
+
+        # TODO: embedding should be the seperately defined, so refactor this part later
+        if llm == 'openai':
+            # self.llm = LangChainGPT()
+            self.llm, self.tokenizer = self.get_models('openai')
+        elif llm == 'debug':
+            self.llm, self.tokenizer = self.get_models( 'debug')
+        elif llm == 'spark':
+            self.llm, self.tokenizer = self.get_models( 'spark')
+        elif llm == 'GLMPro':
+            self.llm, self.tokenizer = self.get_models( 'GLMPro')
+        elif llm == 'ChatGLM2GPT':
+            self.llm, self.tokenizer = self.get_models( 'ChatGLM2GPT')
+            self.story_prefix_prompt = '\n'
+        else:
+            print(f'warning! undefined llm {llm}, use openai instead.')
+            self.llm, self.tokenizer = self.get_models('openai')
+
+        if embedding == 'luotuo_openai':
+            self.embedding = luotuo_openai_embedding
+        else:
+            print(f'warning! undefined embedding {embedding}, use luotuo_openai instead.')
+            self.embedding = luotuo_openai_embedding
+        
+        if role_name:
+
+            from .role_name_to_file import get_folder_role_name
+            # correct role_name to folder_role_name
+            role_name, url = get_folder_role_name(role_name)
+
+            unzip_folder = f'./temp_character_folder/temp_{role_name}'
+            db_folder = os.path.join(unzip_folder, f'content/{role_name}')
+            system_prompt = os.path.join(unzip_folder, f'content/system_prompt.txt')
+
+            if not os.path.exists(unzip_folder):
+                # not yet downloaded
+                # url = f'https://github.com/LC1332/Haruhi-2-Dev/raw/main/data/character_in_zip/{role_name}.zip'
+                import requests, zipfile, io
+                r = requests.get(url)
+                z = zipfile.ZipFile(io.BytesIO(r.content))
+                z.extractall(unzip_folder)
+
+            if self.verbose:
+                print(f'loading pre-defined character {role_name}...')
+            
+            self.db = ChromaDB()
+            self.db.load(db_folder)
+            self.system_prompt = self.check_system_prompt(system_prompt)
+
+        elif story_db:
+            self.db = ChromaDB() 
+            self.db.load(story_db)
+        elif story_text_folder:
+            # print("Building story database from texts...")
+            self.db = self.build_story_db(story_text_folder) 
+        else:
+            self.db = None
+            print('warning! database not yet figured out, both story_db and story_text_folder are not inputted.')
+            # raise ValueError("Either story_db or story_text_folder must be provided")
+        
+
+        self.max_len_story, self.max_len_history = self.get_tokenlen_setting('openai')
+
+        if max_len_history is not None:
+            self.max_len_history = max_len_history
+            # user setting will override default setting
+
+        if max_len_story is not None:
+            self.max_len_story = max_len_story
+            # user setting will override default setting
+
+        self.dialogue_history = []
+
+        
+
+    def check_system_prompt(self, system_prompt):
+        # if system_prompt end with .txt, read the file with utf-8
+        # else, return the string directly
+        if system_prompt.endswith('.txt'):
+            with open(system_prompt, 'r', encoding='utf-8') as f:
+                return f.read()
+        else:
+            return system_prompt
+    
+
+    def get_models(self, model_name):
+
+        # TODO: if output only require tokenizer model, no need to initialize llm
+        
+        # return the combination of llm, embedding and tokenizer
+        if model_name == 'openai':
+            from .LangChainGPT import LangChainGPT
+            return (LangChainGPT(), tiktokenizer)
+        elif model_name == 'debug':
+            from .PrintLLM import PrintLLM
+            return (PrintLLM(), tiktokenizer)
+        elif model_name == 'spark':
+            from .SparkGPT import SparkGPT
+            return (SparkGPT(), tiktokenizer)
+        elif model_name == 'GLMPro':
+            from .GLMPro import GLMPro
+            return (GLMPro(), tiktokenizer)
+        elif model_name == "ChatGLM2GPT":
+            from .ChatGLM2GPT import ChatGLM2GPT, GLM_tokenizer
+            return (ChatGLM2GPT(), GLM_tokenizer)
+        else:
+            print(f'warning! undefined model {model_name}, use openai instead.')
+            from .LangChainGPT import LangChainGPT
+            return (LangChainGPT(), tiktokenizer)
+        
+    def get_tokenlen_setting( self, model_name ):
+        # return the setting of story and history token length
+        if model_name == 'openai':
+            return (1500, 1200)
+        else:
+            print(f'warning! undefined model {model_name}, use openai instead.')
+            return (1500, 1200)
+        
+    def build_story_db_from_vec( self, texts, vecs ):
+        self.db = ChromaDB()
+
+        self.db.init_from_docs( vecs, texts)
+
+    def build_story_db(self, text_folder):
+        # 实现读取文本文件夹,抽取向量的逻辑
+        db = ChromaDB()
+
+        strs = []
+
+        # scan all txt file from text_folder
+        for file in os.listdir(text_folder):
+            # if file name end with txt
+            if file.endswith(".txt"):
+                file_path = os.path.join(text_folder, file)
+                with open(file_path, 'r', encoding='utf-8') as f:
+                    strs.append(f.read())
+
+        if self.verbose:
+            print(f'starting extract embedding... for { len(strs) } files')
+
+        vecs = []
+
+        ## TODO: 建立一个新的embedding batch test的单元测试
+        ## 新的支持list batch test的embedding代码
+        ## 用新的代码替换下面的for循环
+        ## Luotuo-bert-en也发布了，所以可以避开使用openai
+        
+        for mystr in strs:
+            vecs.append(self.embedding(mystr))
+
+        db.init_from_docs(vecs, strs)
+
+        return db
+    
+    def save_story_db(self, db_path):
+        self.db.save(db_path)
+        
+    def chat(self, text, role):
+        # add system prompt
+        self.llm.initialize_message()
+        self.llm.system_message(self.system_prompt)
+    
+
+        # add story
+        query = self.get_query_string(text, role)
+        self.add_story( query )
+
+        # add history
+        self.add_history()
+
+        # add query
+        self.llm.user_message(query)
+        
+        # get response
+        response_raw = self.llm.get_response()
+
+        response = response_postprocess(response_raw, self.dialogue_bra_token, self.dialogue_ket_token)
+
+        # record dialogue history
+        self.dialogue_history.append((query, response))
+
+
+
+        return response
+    
+    def get_query_string(self, text, role):
+        if role in self.narrator:
+            return role + ":" + text
+        else:
+            return f"{role}:{self.dialogue_bra_token}{text}{self.dialogue_ket_token}"
+        
+    def add_story(self, query):
+
+        if self.db is None:
+            return
+        
+        query_vec = self.embedding(query)
+
+        stories = self.db.search(query_vec, self.k_search)
+        
+        story_string = self.story_prefix_prompt
+        sum_story_token = self.tokenizer(story_string)
+        
+        for story in stories:
+            story_token = self.tokenizer(story) + self.tokenizer(self.dialogue_divide_token)
+            if sum_story_token + story_token > self.max_len_story:
+                break
+            else:
+                sum_story_token += story_token
+                story_string += story + self.dialogue_divide_token
+
+        self.llm.user_message(story_string)
+        
+    def add_history(self):
+
+        if len(self.dialogue_history) == 0:
+            return
+        
+        sum_history_token = 0
+        flag = 0
+        for query, response in reversed(self.dialogue_history):
+            current_count = 0
+            if query is not None:
+                current_count += self.tokenizer(query) 
+            if response is not None:
+                current_count += self.tokenizer(response)
+            sum_history_token += current_count
+            if sum_history_token > self.max_len_history:
+                break
+            else:
+                flag += 1
+
+        if flag == 0:
+            print('warning! no history added. the last dialogue is too long.')
+
+        for (query, response) in self.dialogue_history[-flag:]:
+            if query is not None:
+                self.llm.user_message(query)
+            if response is not None:
+                self.llm.ai_message(response)
+
+        
+        

ChatHaruhi/ChromaDB.py

@@ -0,0 +1,61 @@
+import chromadb
+from .BaseDB import BaseDB
+import random
+import string
+import os
+
+class ChromaDB(BaseDB):
+    
+    def __init__(self):
+        self.client = None
+        self.collection = None
+        self.path = None
+    
+    def init_db(self):
+
+        if self.client is not None:
+            print('ChromaDB has already been initialized')
+            return
+
+        folder_name = ''
+
+        while os.path.exists(folder_name) or folder_name == '':
+            # try to create a folder named temp_<random string> which is not yet existed
+            folder_name =  "tempdb_" + ''.join(random.sample(string.ascii_letters + string.digits, 8))
+
+        self.path = folder_name
+        self.client = chromadb.PersistentClient(path = folder_name)
+
+        self.collection = self.client.get_or_create_collection("search")
+
+    def save(self, file_path):
+        if file_path != self.path:
+            # copy all files in self.path to file_path, with overwrite
+            os.system("cp -r " + self.path + " " + file_path)
+            previous_path = self.path
+            self.path = file_path
+            self.client = chromadb.PersistentClient(path = file_path)
+            # remove previous path if it start with tempdb
+            if previous_path.startswith("tempdb"):
+                os.system("rm -rf " + previous_path)
+                        
+
+    def load(self, file_path):
+        self.path = file_path
+        self.client = chromadb.PersistentClient(path = file_path)
+        self.collection = self.client.get_collection("search")
+
+    def search(self, vector, n_results):
+        results = self.collection.query(query_embeddings=[vector], n_results=n_results)
+        return results['documents'][0]
+
+    def init_from_docs(self, vectors, documents):
+        if self.client is None:
+            self.init_db()
+        
+        ids = []
+        for i, doc in enumerate(documents):
+            first_four_chat = doc[:min(4, len(doc))]
+            ids.append( str(i) + "_" + doc)
+        self.collection.add(embeddings=vectors, documents=documents, ids = ids)
+        

ChatHaruhi/GLMPro.py

@@ -0,0 +1,90 @@
+from .BaseLLM import BaseLLM
+import os
+
+zhipu_api = os.environ['ZHIPU_API']
+
+import zhipuai
+import time
+
+class GLMPro( BaseLLM ):
+    def __init__(self, model="chatglm_pro", verbose = False ):
+        super(GLMPro,self).__init__()
+
+        zhipuai.api_key = zhipu_api
+
+        self.verbose = verbose
+
+        self.model_name = model
+
+        self.prompts = []
+
+        if self.verbose == True:
+            print('model name, ', self.model_name )
+            if len( zhipu_api ) > 8:
+                print( 'found apikey ', zhipu_api[:4], '****', zhipu_api[-4:] )
+            else:
+                print( 'found apikey but too short, ' )
+        
+
+    def initialize_message(self):
+        self.prompts = []
+
+    def ai_message(self, payload):
+        self.prompts.append({"role":"assistant","content":payload})
+
+    def system_message(self, payload):
+        self.prompts.append({"role":"user","content":payload})
+
+    def user_message(self, payload):
+        self.prompts.append({"role":"user","content":payload})
+
+    def get_response(self):
+        zhipuai.api_key = zhipu_api
+        max_test_name = 5
+        sleep_interval = 3
+
+        request_id = None
+
+        
+
+        # try submit asychonize request until success
+        for test_time in range( max_test_name ):
+            response = zhipuai.model_api.async_invoke(
+                model = self.model_name,
+                prompt = self.prompts,
+                temperature = 0)
+            if response['success'] == True:
+                request_id = response['data']['task_id']
+
+                if self.verbose == True:
+                    print('submit request, id = ', request_id )
+                break
+            else:
+                print('submit GLM request failed, retrying...')
+                time.sleep( sleep_interval )
+
+        if request_id:
+            # try get response until success
+            for test_time in range( 2 * max_test_name ):
+                result = zhipuai.model_api.query_async_invoke_result( request_id )
+                if result['code'] == 200 and result['data']['task_status'] == 'SUCCESS':
+
+                    if self.verbose == True:
+                        print('get GLM response success' )
+
+                    choices = result['data']['choices']
+                    if len( choices ) > 0:
+                        return choices[-1]['content'].strip("\"'")
+                    
+                # other wise means failed
+                if self.verbose == True:
+                    print('get GLM response failed, retrying...')
+                # sleep for 1 second
+                time.sleep( sleep_interval )
+        else:
+            print('submit GLM request failed, please check your api key and model name')
+            return ''
+    
+    def print_prompt(self):
+        for message in self.prompts:
+            print(f"{message['role']}: {message['content']}")

ChatHaruhi/LangChainGPT.py

@@ -0,0 +1,66 @@
+# ChatHaruhi: Reviving Anime Character in Reality via Large Language Model
+#
+# ChatHaruhi 2.0, built by Cheng Li and Weishi Mi
+#
+# [email protected], [email protected]
+# 
+# Weishi Mi is a second-year graduate student at Tsinghua University, majoring in computer science.
+# Weishi Mi is pursuing a job or a PhD position, which who will be available next year
+# 
+# homepage https://github.com/LC1332/Chat-Haruhi-Suzumiya
+# 
+# ChatHaruhi is a chatbot that can revive anime characters in reality.
+# the 2.0 version was built by Cheng Li and Weishi Mi.
+# 
+# Please cite our paper if you use this code for research: 
+#
+# @misc{li2023chatharuhi,
+#       title={ChatHaruhi: Reviving Anime Character in Reality via Large Language Model}, 
+#       author={Cheng Li and Ziang Leng and Chenxi Yan and Junyi Shen and Hao Wang and Weishi MI and Yaying Fei and Xiaoyang Feng and Song Yan and HaoSheng Wang and Linkang Zhan and Yaokai Jia and Pingyu Wu and Haozhen Sun},
+#       year={2023},
+#       eprint={2308.09597},
+#       archivePrefix={arXiv},
+#       primaryClass={cs.CL}
+# }
+
+
+from langchain.chat_models import ChatOpenAI
+from langchain.prompts.chat import (
+    ChatPromptTemplate,
+    SystemMessagePromptTemplate,
+    AIMessagePromptTemplate,
+    HumanMessagePromptTemplate,
+)
+from langchain.schema import (
+    AIMessage,
+    HumanMessage,
+    SystemMessage
+)
+from .BaseLLM import BaseLLM
+
+class LangChainGPT(BaseLLM):
+
+    def __init__(self, model="gpt-3.5-turbo"):
+        super(LangChainGPT,self).__init__()
+        self.chat = ChatOpenAI(model=model)
+        self.messages = []
+
+    def initialize_message(self):
+        self.messages = []
+
+    def ai_message(self, payload):
+        self.messages.append(AIMessage(content = payload))
+
+    def system_message(self, payload):
+        self.messages.append(SystemMessage(content = payload))
+
+    def user_message(self, payload):
+        self.messages.append(HumanMessage(content = payload))
+
+    def get_response(self):
+        response = self.chat(self.messages)
+        return response.content
+    
+    def print_prompt(self):
+        for message in self.messages:
+            print(message)

ChatHaruhi/PrintLLM.py

@@ -0,0 +1,61 @@
+# ChatHaruhi: Reviving Anime Character in Reality via Large Language Model
+#
+# ChatHaruhi 2.0, built by Cheng Li and Weishi Mi
+#
+# [email protected], [email protected]
+# 
+# Weishi Mi is a second-year graduate student at Tsinghua University, majoring in computer science.
+# Weishi Mi is pursuing a job or a PhD position, which who will be available next year
+# 
+# homepage https://github.com/LC1332/Chat-Haruhi-Suzumiya
+# 
+# ChatHaruhi is a chatbot that can revive anime characters in reality.
+# the 2.0 version was built by Cheng Li and Weishi Mi.
+# 
+# Please cite our paper if you use this code for research: 
+#
+# @misc{li2023chatharuhi,
+#       title={ChatHaruhi: Reviving Anime Character in Reality via Large Language Model}, 
+#       author={Cheng Li and Ziang Leng and Chenxi Yan and Junyi Shen and Hao Wang and Weishi MI and Yaying Fei and Xiaoyang Feng and Song Yan and HaoSheng Wang and Linkang Zhan and Yaokai Jia and Pingyu Wu and Haozhen Sun},
+#       year={2023},
+#       eprint={2308.09597},
+#       archivePrefix={arXiv},
+#       primaryClass={cs.CL}
+# }
+# 
+# This PrintLLM.py is for debuging with any real-runing LLM
+# so you can see full prompt and copy it into GPT or Claude to debug
+#
+
+from .BaseLLM import BaseLLM
+
+class PrintLLM(BaseLLM):
+
+    def __init__(self ):
+        self.messages = []
+        self.messages.append("Noticing: This is a print LLM for debug.")
+        self.messages.append("But you can also copy the prompt into GPT or Claude to debugging")
+
+    def initialize_message(self):
+        self.messages = []
+        self.messages.append("Noticing: This is a print LLM for debug.")
+        self.messages.append("But you can also copy the prompt into GPT or Claude to debugging")
+
+    def ai_message(self, payload):
+        self.messages.append("AI: \n" + payload)
+
+    def system_message(self, payload):
+        self.messages.append("System: \n" + payload)
+
+    def user_message(self, payload):
+        self.messages.append("User: \n" + payload)
+
+    def get_response(self):
+        for message in self.messages:
+            print(message)
+        response = input("Please input your response: ")
+        return response
+    
+    def print_prompt(self):
+        for message in self.messages:
+            print(message)

ChatHaruhi/SparkApi.py

@@ -0,0 +1,139 @@
+# 由讯飞提供的websocket接口，用于与星火机器人进行交互
+
+import _thread as thread
+import base64
+import datetime
+import hashlib
+import hmac
+import json
+from urllib.parse import urlparse
+import ssl
+from datetime import datetime
+from time import mktime
+from urllib.parse import urlencode
+from wsgiref.handlers import format_date_time
+
+import websocket  # 使用websocket_client
+answer = ""
+
+class Ws_Param(object):
+    # 初始化
+    def __init__(self, APPID, APIKey, APISecret, Spark_url):
+        self.APPID = APPID
+        self.APIKey = APIKey
+        self.APISecret = APISecret
+        self.host = urlparse(Spark_url).netloc
+        self.path = urlparse(Spark_url).path
+        self.Spark_url = Spark_url
+
+    # 生成url
+    def create_url(self):
+        # 生成RFC1123格式的时间戳
+        now = datetime.now()
+        date = format_date_time(mktime(now.timetuple()))
+
+        # 拼接字符串
+        signature_origin = "host: " + self.host + "\n"
+        signature_origin += "date: " + date + "\n"
+        signature_origin += "GET " + self.path + " HTTP/1.1"
+
+        # 进行hmac-sha256进行加密
+        signature_sha = hmac.new(self.APISecret.encode('utf-8'), signature_origin.encode('utf-8'),
+                                 digestmod=hashlib.sha256).digest()
+
+        signature_sha_base64 = base64.b64encode(signature_sha).decode(encoding='utf-8')
+
+        authorization_origin = f'api_key="{self.APIKey}", algorithm="hmac-sha256", headers="host date request-line", signature="{signature_sha_base64}"'
+
+        authorization = base64.b64encode(authorization_origin.encode('utf-8')).decode(encoding='utf-8')
+
+        # 将请求的鉴权参数组合为字典
+        v = {
+            "authorization": authorization,
+            "date": date,
+            "host": self.host
+        }
+        # 拼接鉴权参数，生成url
+        url = self.Spark_url + '?' + urlencode(v)
+        # 此处打印出建立连接时候的url,参考本demo的时候可取消上方打印的注释，比对相同参数时生成的url与自己代码生成的url是否一致
+        return url
+
+
+# 收到websocket错误的处理
+def on_error(ws, error):
+    print("### error:", error)
+
+
+# 收到websocket关闭的处理
+def on_close(ws,one,two):
+    print(" ")
+
+
+# 收到websocket连接建立的处理
+def on_open(ws):
+    thread.start_new_thread(run, (ws,))
+
+
+def run(ws, *args):
+    data = json.dumps(gen_params(appid=ws.appid, domain= ws.domain,question=ws.question))
+    ws.send(data)
+
+
+# 收到websocket消息的处理
+def on_message(ws, message):
+    # print(message)
+    data = json.loads(message)
+    code = data['header']['code']
+    if code != 0:
+        print(f'请求错误: {code}, {data}')
+        ws.close()
+    else:
+        choices = data["payload"]["choices"]
+        status = choices["status"]
+        content = choices["text"][0]["content"]
+        # print(content,end ="")
+        global answer
+        answer += content
+        # print(1)
+        if status == 2:
+            ws.close()
+
+
+def gen_params(appid, domain,question):
+    """
+    通过appid和用户的提问来生成请参数
+    """
+    data = {
+        "header": {
+            "app_id": appid,
+            "uid": "1234"
+        },
+        "parameter": {
+            "chat": {
+                "domain": domain,
+                "random_threshold": 0.5,
+                "max_tokens": 2048,
+                "auditing": "default"
+            }
+        },
+        "payload": {
+            "message": {
+                "text": question
+            }
+        }
+    }
+    return data
+
+
+def main(appid, api_key, api_secret, Spark_url,domain, question):
+    # print("星火:")
+    wsParam = Ws_Param(appid, api_key, api_secret, Spark_url)
+    websocket.enableTrace(False)
+    wsUrl = wsParam.create_url()
+    ws = websocket.WebSocketApp(wsUrl, on_message=on_message, on_error=on_error, on_close=on_close, on_open=on_open)
+    ws.appid = appid
+    ws.question = question
+    ws.domain = domain
+    ws.run_forever(sslopt={"cert_reqs": ssl.CERT_NONE})
+
+

ChatHaruhi/SparkGPT.py

@@ -0,0 +1,54 @@
+# SparkGPT.py
+from . import SparkApi
+#以下密钥信息从os环境获取
+import os
+
+appid = os.environ['APPID']
+api_secret = os.environ['APISecret'] 
+api_key = os.environ['APIKey']
+
+
+from .BaseLLM import BaseLLM
+
+    
+
+
+class SparkGPT(BaseLLM):
+
+    def __init__(self, model="Spark2.0"):
+        super(SparkGPT,self).__init__()
+        if model == "Spark2.0":
+            self.domain = "generalv2"    # v2.0版本
+            self.Spark_url = "ws://spark-api.xf-yun.com/v2.1/chat"  # v2.0环境的地址
+        elif model == "Spark1.5":
+            self.domain = "general"   # v1.5版本
+            self.Spark_url = "ws://spark-api.xf-yun.com/v1.1/chat"  # v1.5环境的地址
+        else:
+            raise Exception("Unknown Spark model")
+        # SparkApi.answer =""
+        self.messages = ''
+        
+
+    def initialize_message(self):
+        self.messages = ''
+
+    def ai_message(self, payload):
+        self.messages = self.messages + "AI: " + payload 
+
+    def system_message(self, payload):
+        self.messages = self.messages + "System: " + payload 
+
+    def user_message(self, payload):
+        self.messages = self.messages + "User: " + payload 
+
+    def get_response(self):
+        # question = checklen(getText("user",Input))
+
+        message_json = [{"role": "user", "content": self.messages}]
+        SparkApi.answer =""
+        SparkApi.main(appid,api_key,api_secret,self.Spark_url,self.domain,message_json)
+        return SparkApi.answer
+    
+    def print_prompt(self):
+        print(type(self.messages))
+        print(self.messages)

ChatHaruhi/__init__.py

@@ -0,0 +1,26 @@
+# ChatHaruhi: Reviving Anime Character in Reality via Large Language Model
+#
+# ChatHaruhi 2.0, built by Cheng Li and Weishi Mi
+#
+# [email protected], [email protected]
+# 
+# Weishi Mi is a second-year graduate student at Tsinghua University, majoring in computer science.
+# Weishi Mi is pursuing a job or a PhD position, which who will be available next year
+# 
+# homepage https://github.com/LC1332/Chat-Haruhi-Suzumiya
+# 
+# ChatHaruhi is a chatbot that can revive anime characters in reality.
+# the 2.0 version was built by Cheng Li and Weishi Mi.
+# 
+# Please cite our paper if you use this code for research: 
+#
+# @misc{li2023chatharuhi,
+#       title={ChatHaruhi: Reviving Anime Character in Reality via Large Language Model}, 
+#       author={Cheng Li and Ziang Leng and Chenxi Yan and Junyi Shen and Hao Wang and Weishi MI and Yaying Fei and Xiaoyang Feng and Song Yan and HaoSheng Wang and Linkang Zhan and Yaokai Jia and Pingyu Wu and Haozhen Sun},
+#       year={2023},
+#       eprint={2308.09597},
+#       archivePrefix={arXiv},
+#       primaryClass={cs.CL}
+# }
+
+from .ChatHaruhi import ChatHaruhi

ChatHaruhi/role_name_to_file.py

@@ -0,0 +1,67 @@
+# ChatHaruhi: Reviving Anime Character in Reality via Large Language Model
+#
+# ChatHaruhi 2.0, built by Cheng Li and Weishi Mi
+#
+# [email protected], [email protected]
+# 
+# Weishi Mi is a second-year graduate student at Tsinghua University, majoring in computer science.
+# Weishi Mi is pursuing a job or a PhD position, which who will be available next year
+# 
+# homepage https://github.com/LC1332/Chat-Haruhi-Suzumiya
+# 
+# ChatHaruhi is a chatbot that can revive anime characters in reality.
+# the 2.0 version was built by Cheng Li and Weishi Mi.
+# 
+# Please cite our paper if you use this code for research: 
+#
+# @misc{li2023chatharuhi,
+#       title={ChatHaruhi: Reviving Anime Character in Reality via Large Language Model}, 
+#       author={Cheng Li and Ziang Leng and Chenxi Yan and Junyi Shen and Hao Wang and Weishi MI and Yaying Fei and Xiaoyang Feng and Song Yan and HaoSheng Wang and Linkang Zhan and Yaokai Jia and Pingyu Wu and Haozhen Sun},
+#       year={2023},
+#       eprint={2308.09597},
+#       archivePrefix={arXiv},
+#       primaryClass={cs.CL}
+# }
+# 
+# if you have attempt to add a new character, please add the role name here
+# 
+
+role_name_Haruhiu = {'汤师爷': 'tangshiye', 'tangshiye': 'tangshiye', 'Tangshiye': 'tangshiye', 
+                     '慕容复': 'murongfu', 'murongfu': 'murongfu', 'Murongfu': 'murongfu', 
+                     '李云龙': 'liyunlong', 'liyunlong': 'liyunlong', 'Liyunlong': 'liyunlong', 
+                     'Luna': 'Luna', '王多鱼': 'wangduoyu', 'wangduoyu': 'wangduoyu', 
+                     'Wangduoyu': 'wangduoyu', 'Ron': 'Ron', '鸠摩智': 'jiumozhi', 
+                     'jiumozhi': 'jiumozhi', 'Jiumozhi': 'jiumozhi', 'Snape': 'Snape', 
+                     '凉宫春日': 'haruhi', 'haruhi': 'haruhi', 'Haruhi': 'haruhi', 
+                     'Malfoy': 'Malfoy', '虚竹': 'xuzhu', 'xuzhu': 'xuzhu', 
+                     'Xuzhu': 'xuzhu', '萧峰': 'xiaofeng', 
+                     'xiaofeng': 'xiaofeng', 'Xiaofeng': 'xiaofeng', '段誉': 'duanyu', 
+                     'duanyu': 'duanyu', 'Duanyu': 'duanyu', 'Hermione': 'Hermione', 
+                     'Dumbledore': 'Dumbledore', '王语嫣': 'wangyuyan', 'wangyuyan': 
+                     'wangyuyan', 'Wangyuyan': 'wangyuyan', 'Harry': 'Harry', 
+                     'McGonagall': 'McGonagall', '白展堂': 'baizhantang', 
+                     'baizhantang': 'baizhantang', 'Baizhantang': 'baizhantang', 
+                     '佟湘玉': 'tongxiangyu', 'tongxiangyu': 'tongxiangyu', 
+                     'Tongxiangyu': 'tongxiangyu', '郭芙蓉': 'guofurong', 
+                     'guofurong': 'guofurong', 'Guofurong': 'guofurong', '流浪者': 'wanderer', 
+                     'wanderer': 'wanderer', 'Wanderer': 'wanderer', '钟离': 'zhongli', 
+                     'zhongli': 'zhongli', 'Zhongli': 'zhongli', '胡桃': 'hutao', 'hutao': 'hutao', 
+                     'Hutao': 'hutao', 'Sheldon': 'Sheldon', 'Raj': 'Raj', 
+                     'Penny': 'Penny', '韦小宝': 'weixiaobao', 'weixiaobao': 'weixiaobao', 
+                     'Weixiaobao': 'weixiaobao', '乔峰': 'qiaofeng', 'qiaofeng': 'qiaofeng', 
+                     'Qiaofeng': 'qiaofeng', '神里绫华': 'ayaka', 'ayaka': 'ayaka', 
+                     'Ayaka': 'ayaka', '雷电将军': 'raidenShogun', 'raidenShogun': 'raidenShogun', 
+                     'RaidenShogun': 'raidenShogun', '于谦': 'yuqian', 'yuqian': 'yuqian', 
+                     'Yuqian': 'yuqian', 'Professor McGonagall': 'McGonagall', 
+                     'Professor Dumbledore': 'Dumbledore'}
+
+# input role_name , nick name is also allowed
+# output folder_role_name and url url = f'https://github.com/LC1332/Haruhi-2-Dev/raw/main/data/character_in_zip/{role_name}.zip'
+def get_folder_role_name(role_name):
+    if role_name in role_name_Haruhiu:
+        folder_role_name = role_name_Haruhiu[role_name]
+        url = f'https://github.com/LC1332/Haruhi-2-Dev/raw/main/data/character_in_zip/{folder_role_name}.zip'
+        return folder_role_name, url
+    else:
+        print('role_name {} not found, using haruhi as default'.format(role_name))
+        return get_folder_role_name('haruhi')

ChatHaruhi/utils.py

@@ -0,0 +1,180 @@
+from argparse import Namespace
+
+import openai
+from transformers import AutoModel, AutoTokenizer
+import torch
+import random
+
+import tiktoken
+import re
+
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+_luotuo_model = None
+
+_luotuo_model_en = None
+_luotuo_en_tokenizer = None
+
+_enc_model = None
+
+def tiktokenizer( text ):
+    global _enc_model
+
+    if _enc_model is None:
+        _enc_model = tiktoken.get_encoding("cl100k_base")
+
+    return len(_enc_model.encode(text))
+    
+def response_postprocess(text,dialogue_bra_token = '「',dialogue_ket_token = '」'):
+    lines = text.split('\n')
+    new_lines = ""
+
+    first_name = None
+
+    for line in lines:
+        line = line.strip(" ")
+        match = re.match(r'^(.*?)[:：]' + dialogue_bra_token + r"(.*?)" + dialogue_ket_token + r"$", line)
+
+        
+        if match:
+            curr_name = match.group(1)
+            # print(curr_name)
+            if first_name is None:
+                first_name = curr_name
+                new_lines += (match.group(2))
+            else:
+                if curr_name != first_name:
+                    return first_name + ":" + dialogue_bra_token +  new_lines + dialogue_ket_token
+                else:
+                    new_lines += (match.group(2))
+            
+        else:
+            if first_name == None:
+                return text
+            else:
+                return first_name + ":" + dialogue_bra_token +  new_lines + dialogue_ket_token
+    return first_name + ":" + dialogue_bra_token + new_lines + dialogue_ket_token
+
+def download_models():
+    print("正在下载Luotuo-Bert")
+    # Import our models. The package will take care of downloading the models automatically
+    model_args = Namespace(do_mlm=None, pooler_type="cls", temp=0.05, mlp_only_train=False,
+                           init_embeddings_model=None)
+    model = AutoModel.from_pretrained("silk-road/luotuo-bert-medium", trust_remote_code=True, model_args=model_args).to(
+        device)
+    print("Luotuo-Bert下载完毕")
+    return model
+
+def get_luotuo_model():
+    global _luotuo_model
+    if _luotuo_model is None:
+        _luotuo_model = download_models()
+    return _luotuo_model
+
+
+def luotuo_embedding(model, texts):
+    # Tokenize the texts_source
+    tokenizer = AutoTokenizer.from_pretrained("silk-road/luotuo-bert-medium")
+    inputs = tokenizer(texts, padding=True, truncation=False, return_tensors="pt")
+    inputs = inputs.to(device)
+    # Extract the embeddings
+    # Get the embeddings
+    with torch.no_grad():
+        embeddings = model(**inputs, output_hidden_states=True, return_dict=True, sent_emb=True).pooler_output
+    return embeddings
+
+def luotuo_en_embedding( texts ):
+    # this function implemented by Cheng
+    global _luotuo_model_en
+    global _luotuo_en_tokenizer
+
+    if _luotuo_model_en is None:
+        _luotuo_en_tokenizer = AutoTokenizer.from_pretrained("silk-road/luotuo-bert-en")
+        _luotuo_model_en = AutoModel.from_pretrained("silk-road/luotuo-bert-en").to(device)
+
+    if _luotuo_en_tokenizer is None:
+        _luotuo_en_tokenizer = AutoTokenizer.from_pretrained("silk-road/luotuo-bert-en")
+
+    inputs = _luotuo_en_tokenizer(texts, padding=True, truncation=False, return_tensors="pt")
+    inputs = inputs.to(device)
+
+    with torch.no_grad():
+        embeddings = _luotuo_model_en(**inputs, output_hidden_states=True, return_dict=True, sent_emb=True).pooler_output
+        
+    return embeddings
+
+
+def get_embedding_for_chinese(model, texts):
+    model = model.to(device)
+    # str or strList
+    texts = texts if isinstance(texts, list) else [texts]
+    # 截断
+    for i in range(len(texts)):
+        if len(texts[i]) > 510:
+            texts[i] = texts[i][:510]
+    if len(texts) >= 64:
+        embeddings = []
+        chunk_size = 64
+        for i in range(0, len(texts), chunk_size):
+            embeddings.append(luotuo_embedding(model, texts[i: i + chunk_size]))
+        return torch.cat(embeddings, dim=0)
+    else:
+        return luotuo_embedding(model, texts)
+
+
+def is_chinese_or_english(text):
+    text = list(text)
+    is_chinese, is_english = 0, 0
+
+    for char in text:
+        # 判断字符的Unicode值是否在中文字符的Unicode范围内
+        if '\u4e00' <= char <= '\u9fa5':
+            is_chinese += 4
+        # 判断字符是否为英文字符（包括大小写字母和常见标点符号）
+        elif ('\u0041' <= char <= '\u005a') or ('\u0061' <= char <= '\u007a'):
+            is_english += 1
+    if is_chinese >= is_english:
+        return "chinese"
+    else:
+        return "english"
+
+
+def get_embedding_for_english(text, model="text-embedding-ada-002"):
+    text = text.replace("\n", " ")
+    return openai.Embedding.create(input=[text], model=model)['data'][0]['embedding']
+
+import os
+
+def luotuo_openai_embedding(texts, is_chinese= None ):
+    """
+        when input is chinese, use luotuo_embedding
+        when input is english, use openai_embedding
+        texts can be a list or a string
+        when texts is a list, return a list of embeddings, using batch inference
+        when texts is a string, return a single embedding
+    """
+
+    openai_key = os.environ.get("OPENAI_API_KEY")
+
+    if isinstance(texts, list):
+        index = random.randint(0, len(texts) - 1)
+        if openai_key is None or is_chinese_or_english(texts[index]) == "chinese":
+            return [embed.cpu().tolist() for embed in get_embedding_for_chinese(get_luotuo_model(), texts)]
+        else:
+            return [get_embedding_for_english(text) for text in texts]
+    else:
+        if openai_key is None or is_chinese_or_english(texts) == "chinese":
+            return get_embedding_for_chinese(get_luotuo_model(), texts)[0].cpu().tolist()
+        else:
+            return get_embedding_for_english(texts)
+
+
+# compute cosine similarity between two vector
+def get_cosine_similarity( v1, v2):
+    v1 = torch.tensor(v1).to(device)
+    v2 = torch.tensor(v2).to(device)
+    return torch.cosine_similarity(v1, v2, dim=0).item()
+
+    
+