add translate

app.py

@@ -4,7 +4,7 @@ import requests
 import pytz
 import yaml
 from tools.final_answer import FinalAnswerTool
-
+from tools.translate import TranslateTool
 from Gradio_UI import GradioUI
 
 # Below is an example of a tool that does nothing. Amaze us with your creativity !
@@ -47,15 +47,16 @@ custom_role_conversions=None,
 )
 
 
-# Import tool from Hub
-image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
+# Import tooli from Hub
+image_generaton_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
+translate_tool = TranslateTool()
 
 with open("prompts.yaml", 'r') as stream:
     prompt_templates = yaml.safe_load(stream)
     
 agent = CodeAgent(
     model=model,
-    tools=[final_answer], ## add your tools here (don't remove final answer)
+    tools=[final_answer, image_generaton_tool, translate_tool], ## add your tools here (don't remove final answer)
     max_steps=6,
     verbosity_level=1,
     grammar=None,

tools/translate.py

@@ -0,0 +1,19 @@
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from smolagents.tools import Tool
+
+class TranslateTool(Tool):
+    name = "translate"
+    description = "Translate text from Chinese to English"
+    inputs = {"text": {"type": "string", "description": "The text to translate"}}
+    output_type = "string"
+
+    def forward(self, text: str) -> str:
+        inputs = self.tokenizer(text, return_tensors="pt")
+        outputs = self.model.generate(**inputs, max_new_tokens=512)
+        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)   
+        self.model_id = "ModelSpace/GemmaX2-28-9B-v0.1"
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)
+        self.model = AutoModelForCausalLM.from_pretrained(self.model_id)