|
|
import gradio as gr |
|
|
from gradio import ChatMessage |
|
|
from transformers import ReactCodeAgent, HfApiEngine |
|
|
from utils import stream_from_transformers_agent |
|
|
from prompts import SQUAD_REACT_CODE_SYSTEM_PROMPT |
|
|
from tools.squad_retriever import SquadRetrieverTool |
|
|
from tools.text_to_image import TextToImageTool |
|
|
from dotenv import load_dotenv |
|
|
|
|
|
load_dotenv() |
|
|
|
|
|
TASK_SOLVING_TOOLBOX = [ |
|
|
SquadRetrieverTool(), |
|
|
TextToImageTool(), |
|
|
] |
|
|
|
|
|
model_name = "meta-llama/Meta-Llama-3.1-8B-Instruct" |
|
|
|
|
|
|
|
|
llm_engine = HfApiEngine(model_name) |
|
|
|
|
|
|
|
|
agent = ReactCodeAgent( |
|
|
tools=TASK_SOLVING_TOOLBOX, |
|
|
llm_engine=llm_engine, |
|
|
system_prompt=SQUAD_REACT_CODE_SYSTEM_PROMPT, |
|
|
) |
|
|
|
|
|
def append_example_message(x: gr.SelectData, messages): |
|
|
if x.value["text"] is not None: |
|
|
message = x.value["text"] |
|
|
if "files" in x.value: |
|
|
if isinstance(x.value["files"], list): |
|
|
message = "Here are the files: " |
|
|
for file in x.value["files"]: |
|
|
message += f"{file}, " |
|
|
else: |
|
|
message = x.value["files"] |
|
|
messages.append(ChatMessage(role="user", content=message)) |
|
|
return messages |
|
|
|
|
|
def add_message(message, messages): |
|
|
messages.append(ChatMessage(role="user", content=message)) |
|
|
return messages |
|
|
|
|
|
def interact_with_agent(messages): |
|
|
prompt = messages[-1]['content'] |
|
|
for msg in stream_from_transformers_agent(agent, prompt): |
|
|
messages.append(msg) |
|
|
yield messages |
|
|
yield messages |
|
|
|
|
|
with gr.Blocks(fill_height=True) as demo: |
|
|
chatbot = gr.Chatbot( |
|
|
label="SQuAD Agent", |
|
|
type="messages", |
|
|
avatar_images=( |
|
|
None, |
|
|
"https://em-content.zobj.net/source/twitter/53/robot-face_1f916.png", |
|
|
), |
|
|
scale=1, |
|
|
bubble_full_width=False, |
|
|
autoscroll=True, |
|
|
show_copy_all_button=True, |
|
|
show_copy_button=True, |
|
|
placeholder="Enter a message", |
|
|
examples=[ |
|
|
{ |
|
|
"text": "What is on top of the Notre Dame building?", |
|
|
}, |
|
|
{ |
|
|
"text": "Tell me what's on top of the Notre Dame building, and draw a picture of it.", |
|
|
}, |
|
|
{ |
|
|
"text": "Draw a picture of whatever is on top of the Notre Dame building.", |
|
|
}, |
|
|
], |
|
|
) |
|
|
text_input = gr.Textbox(lines=1, label="Chat Message", scale=0) |
|
|
chat_msg = text_input.submit(add_message, [text_input, chatbot], [chatbot]) |
|
|
bot_msg = chat_msg.then(interact_with_agent, [chatbot], [chatbot]) |
|
|
text_input.submit(lambda: "", None, text_input) |
|
|
chatbot.example_select(append_example_message, [chatbot], [chatbot]).then( |
|
|
interact_with_agent, [chatbot], [chatbot] |
|
|
) |
|
|
|
|
|
if __name__ == "__main__": |
|
|
demo.launch() |
|
|
|
