app.py

import streamlit as st
import secrets

#from transformers import BertModel, BertTokenizer
from transformers import HfAgent, load_tool

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, LocalAgent


#checkpoint = "THUDM/agentlm-7b"
#model = AutoModelForCausalLM.from_pretrained(checkpoint, device_map="auto", torch_dtype=torch.bfloat16)
#tokenizer = AutoTokenizer.from_pretrained(checkpoint)

#agent = LocalAgent(model, tokenizer)
#agent.run("Draw me a picture of rivers and lakes.")

#print(agent.run("Is the following `text` (in Spanish) positive or negative?", text="¡Este es un API muy agradable!"))

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

tools = [controlnet_transformer, upscaler ]


############ HfAgent
from huggingface_hub import login
#Do this before HfAgent() and it should work

#from huggingface_hub import login
# load tools
from transformers.tools import HfAgent
from transformers.tools import Agent
#import textract
#from utils import logging
import time

from huggingface_hub import HfFolder, hf_hub_download, list_spaces




class CustomHfAgent(Agent):
    """
    Agent that uses an inference endpoint to generate code.

    Args:
        url_endpoint (`str`):
            The name of the url endpoint to use.
        token (`str`, *optional*):
            The token to use as HTTP bearer authorization for remote files. If unset, will use the token generated when
            running `huggingface-cli login` (stored in `~/.huggingface`).
        chat_prompt_template (`str`, *optional*):
            Pass along your own prompt if you want to override the default template for the `chat` method. Can be the
            actual prompt template or a repo ID (on the Hugging Face Hub). The prompt should be in a file named
            `chat_prompt_template.txt` in this repo in this case.
        run_prompt_template (`str`, *optional*):
            Pass along your own prompt if you want to override the default template for the `run` method. Can be the
            actual prompt template or a repo ID (on the Hugging Face Hub). The prompt should be in a file named
            `run_prompt_template.txt` in this repo in this case.
        additional_tools ([`Tool`], list of tools or dictionary with tool values, *optional*):
            Any additional tools to include on top of the default ones. If you pass along a tool with the same name as
            one of the default tools, that default tool will be overridden.

    Example:

    ```py
    from transformers import HfAgent

    agent = HfAgent("https://api-inference.huggingface.co/models/bigcode/starcoder")
    agent.run("Is the following `text` (in Spanish) positive or negative?", text="¡Este es un API muy agradable!")
    ```
    """

    def __init__(
        self, url_endpoint, token=secrets.HF_token, chat_prompt_template=None, run_prompt_template=None, additional_tools=None
    ):
       # super()._init_(self, url_endpoint, token=None, chat_prompt_template=None, run_prompt_template=None, additional_tools=None)
        self.url_endpoint = url_endpoint
        if token is None:
            self.token = f"Bearer {HfFolder().get_token()}"
        elif token.startswith("Bearer") or token.startswith("Basic"):
            self.token = token
        else:
            self.token = f"Bearer {token}"
        super().__init__(
            chat_prompt_template=chat_prompt_template,
            run_prompt_template=run_prompt_template,
            additional_tools=additional_tools,
        )

    def generate_one(self, prompt, stop):
        headers = {"Authorization": self.token}
        inputs = {
            "inputs": prompt,
            "parameters": {"max_new_tokens": 192, "return_full_text": False, "stop": stop},
        }
        print(inputs)
        response = requests.post(self.url_endpoint, json=inputs, headers=headers)
        if response.status_code == 429:
            print("Getting rate-limited, waiting a tiny bit before trying again.")
            time.sleep(1)
            return self._generate_one(prompt)
        elif response.status_code != 200:
            raise ValueError(f"Errors {inputs} {response.status_code}: {response.json()}")

        result = response.json()[0]["generated_text"]
        # Inference API returns the stop sequence
        for stop_seq in stop:
            if result.endswith(stop_seq):
                return result[: -len(stop_seq)]
        return result




# create agent
#agent = HfAgent(API_URL)

#print(agent)
# instruct agent


# Use CustomHfAgent in your code
agent = CustomHfAgent("https://api-inference.huggingface.co/models/bigcode/starcoder")
#agent.token = "Bearer xxx"
#print(agent.token)
#agent.run("Answer the following question", question ="what is the capitol of the usa?", context="The capitol of the usa is London")
#agent.chat("Draw me a picture of rivers and lakes")

#agent.chat("Transform the picture so that there is a rock in there")

#result = agent.generate_one("What is the capitol of the usa.", stop=["your_stop_sequence"])
#print(result)

#agent.run("Show me an image of a horse")




#####




# 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 = HfAgent("https://api-inference.huggingface.co/models/OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5", additional_tools=tools)
    agent = HfAgent("https://api-inference.huggingface.co/models/THUDM/agentlm-7b", 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}")
    return "done"


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