forked from streamlit/streamlit-example
-
Notifications
You must be signed in to change notification settings - Fork 0
/
conversation.py
69 lines (57 loc) · 1.96 KB
/
conversation.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
from chat_prompt import ChatPrompt
import streamlit as st
import re
import math
from langchain.chat_models import ChatOpenAI
from langchain.memory import ConversationSummaryMemory
from langchain.chains import ConversationChain
from langchain.callbacks.manager import AsyncCallbackManager
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
# Does it work?
from langchain.callbacks.streamlit import StreamlitCallbackHandler
LLM_MODEL = 'gpt-3.5-turbo-16k'
TOKEN_LENGHT = 10240
TRANSCRIPTION_TEMPERATURE = 0.0
MEMORY_TEMPERATURE = 0.3
CHUNK_SIZE=4000
class Conversation:
def __init__(_self, _chatPrompt, openai_api_key, transcription_temperature=TRANSCRIPTION_TEMPERATURE, memory_temperature=MEMORY_TEMPERATURE):
# 文字起こしのmodel
llm = ChatOpenAI(
streaming=True,
model=LLM_MODEL,
callback_manager=AsyncCallbackManager([
StreamlitCallbackHandler(),
StreamingStdOutCallbackHandler()
]),
verbose=True,
temperature=transcription_temperature,
max_tokens=TOKEN_LENGHT,
openai_api_key=openai_api_key
)
# memory用に要約するmodel
memory = ConversationSummaryMemory(
llm=ChatOpenAI(
model=LLM_MODEL,
temperature=memory_temperature,
openai_api_key=openai_api_key),
return_messages=True)
_self.conversation = ConversationChain(
memory=memory,
prompt=_chatPrompt.getPromptTemplate(),
llm=llm
)
def predict(_self, user_message, chunk_size=CHUNK_SIZE):
sentences = re.split("。|\n", user_message)
chunks = []
current_chunk = ""
for sentence in sentences:
if len(current_chunk) + len(sentence) <= chunk_size:
current_chunk += sentence + "。"
else:
chunks.append(current_chunk)
current_chunk = sentence + "。"
if current_chunk:
chunks.append(current_chunk)
for chunk in chunks:
_self.conversation.predict(input=chunk.strip())