-
Notifications
You must be signed in to change notification settings - Fork 2
/
gemini_explorer.py
65 lines (50 loc) · 1.68 KB
/
gemini_explorer.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
import vertexai
import streamlit as st
from vertexai.preview import generative_models
from vertexai.preview.generative_models import GenerativeModel, Part, Content, ChatSession
project = 'radicalx-gemini-explorer'
vertexai.init(project = project)
config = generative_models.GenerationConfig(
temperature = 0.4
)
# Loading the model with the configuration
model = GenerativeModel(
"gemini-pro",
generation_config = config
)
chat = model.start_chat()
# Helper function to display and send streamlit messages
def llm_function(chat: ChatSession, query):
response = chat.send_message(query)
output = response.candidates[0].content.parts[0].text
with st.chat_message("model"):
st.markdown(output)
st.session_state.messages.append({
"role": "user",
"content": query
})
st.session_state.messages.append({
"role": "model",
"content": output
})
st.title("RadicalX Gemini Explorer")
# Initializing the Chat History
if "messages" not in st.session_state:
st.session_state.messages = []
# To display and load the Chat History
for index, message in enumerate(st.session_state.messages):
content = Content(
role = message['role'],
parts = [Part.from_text(message['content'])]
)
chat.history.append(content)
# For initial message startup
if len(st.session_state.messages) == 0:
initial_prompt = f"Introduce yourself as ReX, an AI assistant powered by Geminin Pro. You use emojis to be interactive."
llm_function(chat, initial_prompt)
# To capture the user's input
query = st.chat_input("Message ReX")
if query:
with st.chat_message("user"):
st.markdown(query)
llm_function(chat, query)