/
app.py
70 lines (56 loc) 路 1.83 KB
/
app.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
70
import os
from langchain_community.llms import ollama
from langchain_community.callbacks import StreamlitCallbackHandler
from langchain.prompts import PromptTemplate
# All we need for a good chat
from langchain.memory import ConversationBufferMemory
from langchain.chains import ConversationChain
import streamlit as st
ollama_base_url = os.getenv("OLLAMA_BASE_URL")
memory = ConversationBufferMemory()
## session state variable
if 'chat_history' not in st.session_state:
st.session_state.chat_history=[]
else:
for message in st.session_state.chat_history:
memory.save_context({'input': message['human']}, {'output': message['AI']})
prompt_template = PromptTemplate(
input_variables=['history', 'input'],
template="""
You are a friendly bot.
Conversation history:
{history}
Human: {input}
AI:
"""
)
model = ollama.Ollama(
temperature=0,
repeat_penalty=1,
base_url=ollama_base_url,
model='gemma:2b',
#model='gemma:2b-instruct',
)
conversation_chain = ConversationChain(
prompt=prompt_template,
llm=model,
memory=memory,
verbose=True, # then you can see the intermediate messages
)
# Add a title and a subtitle to the webapp
st.title("馃 I'm Pi-Lot & I 馃挋 Gemma")
st.header("馃憢 I'm running on a PI5")
# Text input fields
user_input = st.chat_input("Topic:")
# Executing the chain when the user
# has entered a topic
if user_input:
st_callback = StreamlitCallbackHandler(st.container())
result = conversation_chain.invoke(
{"input":user_input, "history":st.session_state["chat_history"]},
{"callbacks":[st_callback]}
)
message = {'human': user_input, 'AI': result["response"]}
st.session_state.chat_history.append(message)
with st.expander(label='Chat history', expanded=False):
st.write(st.session_state.chat_history)