import streamlit as st
#from transformers import BertModel, BertTokenizer
from transformers import HfAgent, load_tool

# Load tools
controlnet_transformer = load_tool("huggingface-tools/text-to-image")
upscaler = load_tool("diffusers/latent-upscaler-tool")

tools = [controlnet_transformer, upscaler ]

# Define the model and tokenizer
#model = BertModel.from_pretrained('bert-base-uncased')
#tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

# Create the Streamlit app
st.title("Hugging Face Agent")

# Input field for the user's message
message_input = st.text_input("Enter your message:", "")

# Checkboxes for the tools to be used by the agent
tool_checkboxes = [st.checkbox(f"Use {tool}") for tool in tools]

# Submit button
submit_button = st.button("Submit")

# Define the callback function to handle the form submission
def handle_submission():
    # Get the user's message and the selected tools
    message = message_input
    selected_tools = [tool for tool, checkbox in zip(tools, tool_checkboxes) if checkbox]

    # Initialize the agent with the selected tools
    agent = HfAgent("https://api-inference.huggingface.co/models/bigcode/starcoder", additional_tools=tools)

 #   agent.config.tokenizer = tokenizer
    agent.config.tools = selected_tools

    # Process the user's message
   # inputs = tokenizer.encode_plus(message, add_special_tokens=True, return_tensors="pt")
   # outputs = agent(inputs['input_ids'], attention_mask=inputs['attention_mask'])

    # Display the agent's response
    response = agent.run(message)
    st.text(f"{response:.4f}")

# Add the callback function to the Streamlit app
submit_button = st.button("Submit", on_click=handle_submission)