-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
121 lines (101 loc) · 4.87 KB
/
main.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
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
import streamlit as st
__import__('pysqlite3')
import sys
sys.modules['sqlite3'] = sys.modules.pop('pysqlite3')
import replicate
from langchain_community.vectorstores import Chroma
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import TextLoader
from langchain_community.embeddings.sentence_transformer import (
SentenceTransformerEmbeddings,
)
from langchain.tools.retriever import create_retriever_tool
from langchain_openai import ChatOpenAI
from langchain import hub
from langchain.agents import create_openai_functions_agent
from langchain.agents import AgentExecutor
from langchain_community.tools.tavily_search import TavilySearchResults
import os
from langchain.memory import ConversationBufferMemory
from langchain_community.utilities import GoogleSearchAPIWrapper
from langchain_core.tools import Tool
search = GoogleSearchAPIWrapper()
search_tool = Tool(
name="google_search",
description="Search Google for recent results.",
func=search.run,
)
tools = [search_tool]
# Get the prompt to use - you can modify this!
prompt = hub.pull("hwchase17/openai-functions-agent")
llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0, api_key=st.secrets["api_key"])
agent = create_openai_functions_agent(llm, tools, prompt)
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
# Sidebar to clear chat history
def clear_chat_history():
st.session_state.messages = [{"role": "assistant", "content": "How may I assist you today?"}]
def generate_llava_response(prompt_input):
string_dialogue = "You are a helpful assistant. You do not respond as 'User' or pretend to be 'User'. You only respond once as 'Assistant'."
for dict_message in st.session_state.messages:
if dict_message["role"] == "user":
if dict_message.get("type") == "image":
string_dialogue = "User: [Image]\\n\\n"
else:
string_dialogue = "User: " dict_message["content"] "\\n\\n"
else:
string_dialogue = "Assistant: " dict_message["content"] "\\n\\n"
output = replicate.run(llm,
input={"image": uploaded_file, "prompt": prompt})
return output
# st.set_page_config(page_title="🦙💬 Llava 2 Chatbot")
with (st.sidebar):
st.title('🦙💬 Wazzup!!! Upload Image to assess if it is safe work site.')
# selected_model = st.sidebar.selectbox('Choose a llava model', ['llava-13b'], key='selected_model')
# if selected_model == 'llava-13b':
llm = 'yorickvp/llava-13b:2facb4a474a0462c15041b78b1ad70952ea46b5ec6ad29583c0b29dbd4249591'
temperature = 0.01
top_p = 0.9
max_length = value=512
# temperature = st.sidebar.slider('temperature', min_value=0.01, max_value=5.0, value=0.1, step=0.01)
# top_p = st.sidebar.slider('top_p', min_value=0.01, max_value=1.0, value=0.9, step=0.01)
# max_length = st.sidebar.slider('max_length', min_value=64, max_value=4096, value=512, step=8)
# st.markdown('📖 Learn how to build a llava chatbot [blog](#link-to-blog)!')
uploaded_file = st.file_uploader("Upload an image", type=["jpg", "png", "jpeg"])
st.sidebar.button('Clear Chat History', on_click=clear_chat_history)
# Ensure session state for messages if it doesn't exist
if "messages" not in st.session_state:
st.session_state.messages = [{"role": "assistant", "content": "How may I assist you today?"}]
# Display or clear chat messages
for message in st.session_state.messages:
with st.chat_message(message["role"]):
if message.get("type") == "image":
st.image(message["content"])
else:
st.write(message["content"])
# Text input for chat
# prompt = st.text_input("Type a message:")
prompt = "identify the work health and safety issues at this site."
# Button to send the message/image
if st.button('Identiy WHS Issues'):
if uploaded_file:
# If an image is uploaded, store it in session_state
st.session_state.messages.append({"role": "user", "content": uploaded_file, "type": "image"})
if prompt:
st.session_state.messages.append({"role": "user", "content": prompt})
llava2_answer = generate_llava_response(st.session_state.messages)
st.write(llava2_answer)
result_string = ""
for text in llava2_answer:
result_string = text
result_string = ("Check the Model Work Health and Safety Bill for non conformance based on this "
"text and identify the Part, Division and "
"Subdivision relating to work health and safety "
"non conformance: ") result_string
chat_history = memory.buffer_as_messages
response = agent_executor.invoke({
"input": result_string,
"chat_history": chat_history,
})
# response = str(response)
st.write(response["output"])