Upload tool

app.py

@@ -0,0 +1,6 @@
+from smolagents import launch_gradio_demo
+from tool import TranslationTool
+
+tool = TranslationTool()
+
+launch_gradio_demo(tool)

requirements.txt

@@ -0,0 +1 @@
+smolagents

tool.py

@@ -0,0 +1,43 @@
+from typing import Any, Optional
+from smolagents.tools import Tool
+
+class TranslationTool(Tool):
+    """
+    Example:
+
+    ```py
+    translator = TranslationTool()
+    translator("This is a super nice API!", src_lang="English", tgt_lang="French")
+    ```
+    """
+    default_checkpoint = "facebook/nllb-200-distilled-600M"
+    description = "This is a tool that translates text from a language to another."
+    name = "translator"
+    inputs = {'text': {'type': 'string', 'description': 'The text to translate'}, 'src_lang': {'type': 'string', 'description': "The language of the text to translate. Written in plain English, such as 'Romanian', or 'Albanian'"}, 'tgt_lang': {'type': 'string', 'description': "The language for the desired output language. Written in plain English, such as 'Romanian', or 'Albanian'"}}
+    output_type = "string"
+
+    def __init__(self, lang_to_code=LANGUAGE_CODES, pre_processor_class=AutoTokenizer, model_class=AutoModelForSeq2SeqLM):
+        super().__init__()
+        self.lang_to_code = lang_to_code
+        self.pre_processor_class = pre_processor_class
+        self.model_class = model_class
+        # self.pre_processor = self.pre_processor_class.from_pretrained(self.default_checkpoint)
+        # self.model = self.model_class.from_pretrained(self.default_checkpoint)
+        # self.post_processor = self.pre_processor
+
+    def encode(self, text, src_lang, tgt_lang):
+        if src_lang not in self.lang_to_code:
+            raise ValueError(f"{src_lang} is not a supported language.")
+        if tgt_lang not in self.lang_to_code:
+            raise ValueError(f"{tgt_lang} is not a supported language.")
+        src_lang = self.lang_to_code[src_lang]
+        tgt_lang = self.lang_to_code[tgt_lang]
+        return self.pre_processor._build_translation_inputs(
+            text, return_tensors="pt", src_lang=src_lang, tgt_lang=tgt_lang
+        )
+
+    def forward(self, inputs):
+        return self.model.generate(**inputs)
+
+    def decode(self, outputs):
+        return self.post_processor.decode(outputs[0].tolist(), skip_special_tokens=True)