form chat prompt

model/custom_agent.py

@@ -11,16 +11,28 @@ Dependencies:
 
 Classes:
 - CustomHfAgent: A custom class for interacting with the Hugging Face model API.
+
+Reasono for making this https://github.com/huggingface/transformers/issues/28217 
+Based on https://github.com/huggingface/transformers/blob/main/src/transformers/tools/agents.py
+
+"return_full_text": False,
+
 """
 
 import time
 import requests
 from transformers import Agent
 from utils.logger import log_response
+    
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
 
 class CustomHfAgent(Agent):
     """A custom class for interacting with the Hugging Face model API."""
 
+
+
+
     def __init__(self, url_endpoint, token, chat_prompt_template=None, run_prompt_template=None, additional_tools=None, input_params=None):
         """
         Initialize the CustomHfAgent.
@@ -82,3 +94,49 @@ class CustomHfAgent(Agent):
             if result.endswith(stop_seq):
                 return result[: -len(stop_seq)]
         return result
+###
+    ### 
+    ### https://github.com/huggingface/transformers/blob/main/src/transformers/tools/prompts.py -> run chat_template.txt 
+    ### https://huggingface.co/datasets/huggingface-tools/default-prompts/blob/main/chat_prompt_template.txt
+    ###
+    def format_prompt(self, task, chat_mode=False):
+
+        checkpoint = "bigcode/starcoder"
+        tokenizer = AutoTokenizer.from_pretrained(checkpoint)
+        #model = AutoModelForCausalLM.from_pretrained(checkpoint)  # You may want to use bfloat16 and/or move to GPU here
+
+ 
+
+        description = "\n".join([f"- {name}: {tool.description}" for name, tool in self.toolbox.items()])
+        if chat_mode:
+            if self.chat_history is None:
+                prompt = self.chat_prompt_template.replace("<<all_tools>>", description)
+                messages = [
+                    {
+                        "role": "user",
+                        "content": prompt,
+                    } 
+                ]
+                prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
+            else:
+                prompt = self.chat_history
+                cmp = CHAT_MESSAGE_PROMPT.replace("<<task>>", task)
+                messages = [
+                    {
+                        "role": "user",
+                        "content": cmp,
+                    } 
+                ]
+                cmp = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
+            prompt += cmp
+        else:
+            prompt = self.run_prompt_template.replace("<<all_tools>>", description)
+            prompt = prompt.replace("<<prompt>>", task)
+            messages = [
+                    {
+                        "role": "user",
+                        "content": prompt,
+                    } 
+                ]
+            prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
+        return prompt