Update text_generator.py

text_generator.py

@@ -1,14 +1,20 @@
 import os
 import requests
 import gradio as gr
-from transformers import pipeline, Tool
+from transformers import pipeline
+from smolagents import Tool
 
 class TextGenerationTool(Tool):
     name = "text_generator"
     description = "This is a tool for text generation. It takes a prompt as input and returns the generated text."
     
-    inputs = {"text": {"type": "text", "description": "The prompt for text generation"}}
-    outputs = {"text": {"type": "text", "description": "The generated text"}}
+    inputs = {
+        "text": {
+            "type": "string",
+            "description": "The prompt for text generation"
+        }
+    }
+    output_type = "string"
     
     # Available text generation models
     models = {
@@ -19,8 +25,9 @@ class TextGenerationTool(Tool):
         "openchat": "openchat/openchat_3.5"
     }
     
-    def __init__(self, default_model="orca", use_api=False):
+    def __init__(self, default_model="gpt2", use_api=False):
         """Initialize with a default model and API preference."""
+        super().__init__()
         self.default_model = default_model
         self.use_api = use_api
         self._pipelines = {}
@@ -30,9 +37,9 @@ class TextGenerationTool(Tool):
         if self.token is None and use_api:
             print("Warning: HF_token environment variable not set. API calls will fail.")
     
-    def __call__(self, prompt: str):
+    def forward(self, text: str):
         """Process the input prompt and generate text."""
-        return self.generate_text(prompt)
+        return self.generate_text(text)
     
     def generate_text(self, prompt, model_key=None, max_length=500, temperature=0.7):
         """Generate text based on the prompt using the specified or default model."""