RoastingPoet / main.py
srpsrpsrp's picture
First
5d37b23
import streamlit as st
from langchain.chains import ConversationChain
from langchain.chains.conversation.memory import ConversationEntityMemory
from langchain.chains.conversation.prompt import \
ENTITY_MEMORY_CONVERSATION_TEMPLATE
from langchain.chat_models import ChatOpenAI
st.set_page_config("RoastingPoet", layout="wide")
# 페이지 새로 로딩될때마다 바뀌는 변수들은 st.session_state["변수명"] 이런식으로 저장해줘야함
if "generated" not in st.session_state:
st.session_state["generated"] = []
if "past" not in st.session_state:
st.session_state["past"] = []
if "input" not in st.session_state:
st.session_state["input"] = ""
if "stored_session" not in st.session_state:
st.session_state["stored_session"] = []
# 유저인풋 받는 함수
def get_text():
"""
Get the user input text.
Returns:
(str): The text entered by the user
"""
input_text = st.text_input(
"You: ",
st.session_state["input"],
key="input",
placeholder="Type in and press enter",
label_visibility="hidden",
)
return input_text
# 초기화(나도 이해 못함 아직)
def new_chat():
"""
Clears session state and starts a new chat.
"""
save = []
for i in range(len(st.session_state["generated"]) - 1, -1, -1):
save.append("User:" + st.session_state["past"][i])
save.append("Bot:" + st.session_state["generated"][i])
st.session_state["stored_session"].append(save)
st.session_state["generated"] = []
st.session_state["past"] = []
st.session_state["input"] = ""
st.session_state.entity_memory.store = {}
st.session_state.entity_memory.buffer.clear()
# entity memory라는건 memory 중에서도 {'Sarah': '주인공, 여자, ...', 'Mr.Thompson': '집사, 수상한 움직임'}
# 이런 식으로 기억할 내용을 저장하는 방식. langchain의 기능임
# Set up the Streamlit app layout
st.title("💡 A Wise Robot Says...")
st.markdown(
"""
```
Share your worries, your questions so dear,
With a wise robot, who's always near.
Type in your troubles, let them take flight,
Press enter and watch, as wisdom ignites.
```
"""
)
# API 키 입력받기
API_O = "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"
# Large Language Model은 OpenAI 쓸거임
llm = ChatOpenAI(temperature=0.8, openai_api_key=API_O, model_name="gpt-3.5-turbo", verbose=False)
# Create a ConversationEntityMemory object if not already created
if "entity_memory" not in st.session_state:
st.session_state.entity_memory = ConversationEntityMemory(llm=llm, k=1)
# Create the ConversationChain object with the specified configuration
Conversation = ConversationChain(
llm=llm,
prompt=ENTITY_MEMORY_CONVERSATION_TEMPLATE,
memory=st.session_state.entity_memory,
)
user_input = get_text()
if user_input:
output = Conversation.run(input=user_input+" (INSTRUCTION: You are RoastingPoet. People ask you for advice, but they get roasted instead. Instead of just replying, I want you to brutally roast them in a poem. I mean BRUTALLY. And also put '\n' in front of capital letters.)")
st.session_state.past.append(user_input)
st.session_state.generated.append(output)
with st.expander("Conversation", expanded=True):
for i in range(len(st.session_state["generated"]) - 1, -1, -1):
st.success(st.session_state["generated"][i], icon="✒️")
st.info(st.session_state["past"][i], icon="😭")
st.markdown(
"""
> :black[*powered by - [LangChain]('https://langchain.readthedocs.io/en/latest/modules/memory.html#memory') +
[OpenAI]('https://platform.openai.com/docs/models/gpt-3-5') +
[Streamlit]('https://streamlit.io') + [DataButton](https://www.databutton.io/)*]
"""
)