liuyimeta/training_data_chat

1import pickle
2import faiss
3import openai
4from langchain import LLMChain
5from langchain.llms.openai import OpenAIChat
6from langchain.prompts import Prompt
7from langchain import OpenAI
8from langchain.callbacks import get_openai_callback
9from langchain.callbacks.base import CallbackManager
10from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
11
12
13history = []
14index = 0
15store = None
16prompt = ''
17llmChain = ''
18k = 0
19
20
21def runPrompt(user_input):
22    global index, k, store, prompt, llmChain
23    k += 1
24    if k <= 1:  # 避免重复请求
25        index = faiss.read_index("after_training/training.index")
26        with open("after_training/faiss.pkl", "rb") as f:
27            store = pickle.load(f)
28            store.index = index
29        with open("training/master.txt", "r") as f:
30            promptTemplate = f.read()
31        prompt = Prompt(template=promptTemplate, input_variables=["history", "context", "question"])
32        llmChain = LLMChain(prompt=prompt, llm=OpenAIChat(temperature=0))
33
34    def onMessage(question, history_p):
35        # contexts = []
36        # response_prarm = OpenAI(
37        #     temperature=0,
38        #     openai_api_key=openai.api_key,
39        #     model_name="gpt-3.5-turbo",
40        #     callback_manager=CallbackManager([StreamingStdOutCallbackHandler()]),
41        #     verbose=True,
42        #     streaming=True
43        # )
44        #
45        # llmChain = LLMChain(prompt=prompt, llm=response_prarm)
46        # ai_answer = llmChain.predict(question=question, context="\n\n".join(contexts), history=history_p,
47        #                              stop=["Human:", "AI:"])
48
49        docs = store.similarity_search(question, k=1)
50        contexts = []
51        for i, doc in enumerate(docs):
52            contexts.append(f"Context {i}:\n{doc.page_content}")
53            ai_answer = llmChain.predict(question=question, context="\n\n".join(contexts), history=history_p)
54        return ai_answer
55
56    answer = onMessage(user_input, history)
57    history.append(f"Human: {user_input}")
58    history.append(f"Bot: {answer}")
59    return answer
60
61