Update text_generator.py

text_generator.py

@@ -1,68 +1,93 @@
-import requests
-import os
-from transformers import pipeline
-
-
-from transformers import Tool
-# Import other necessary libraries if needed
-
-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"]
-    outputs = ["text"]
-
-    def __call__(self, prompt: str):
-        #API_URL = "https://api-inference.huggingface.co/models/openchat/openchat_3.5"
-        #headers = {"Authorization": "Bearer " +  os.environ['hf']}
-        token=os.environ['HF_token']
-        #payload = {
-        #    "inputs": prompt  # Adjust this based on your model's input format
-        #}
-
-        #payload = {
-        #            "inputs": "Can you please let us know more details about your ",
-        #        }
-        
-        #def query(payload):
-        #generated_text = requests.post(API_URL, headers=headers, json=payload).json()
-        #print(generated_text)
-        #return generated_text["text"]
-            
-        # Replace the following line with your text generation logic
-        #generated_text = f"Generated text based on the prompt: '{prompt}'"
-
-        # Initialize the text generation pipeline
-        #text_generator = pipeline(model="lgaalves/gpt2-dolly", token=token)
-        text_generator = pipeline(model="microsoft/Orca-2-13b", token=token)
-
-        # Generate text based on a prompt
-        generated_text = text_generator(prompt, max_length=500, num_return_sequences=1, temperature=0.7)
-
-        # Print the generated text
-        print(generated_text)
-
-
-
-        return generated_text
-        
-        # Define the payload for the request
-        #payload = {
-        #    "inputs": prompt  # Adjust this based on your model's input format
-        #}
-
-        # Make the request to the API
-        #generated_text = requests.post(API_URL, headers=headers, json=payload).json()
-
-        # Extract and return the generated text
-        #return generated_text["generated_text"]
-
-# Uncomment and customize the following lines based on your text generation needs
-# text_generator = pipeline(model="gpt2")
-# generated_text = text_generator(prompt, max_length=500, num_return_sequences=1, temperature=0.7)
-
-# Print the generated text if needed
-# print(generated_text)
+import json
+import gradio as gr
+
+def generate_files(title="Text Generation Tool", emoji="🌖", colorFrom="blue", colorTo="blue",
+                   sdk="gradio", sdk_version="4.3.0", app_file="app.py", pinned=False,
+                   tags=["tool"], tool_name="text_generator", tool_description="This is a tool that chats with a user. "
+                   "It takes an input named `prompt` which contains a system_role, user_message, context and history. "
+                   "It returns a text message."):
+    # Generate readme content
+    readme_content = f'''## readme
+title: {title}
+emoji: {emoji}
+colorFrom: {colorFrom}
+colorTo: {colorTo}
+sdk: {sdk}
+sdk_version: {sdk_version}
+app_file: {app_file}
+pinned: {pinned}
+tags:
+  - {tags[0]}
+'''
+
+    # Generate tool config JSON content
+    tool_config = {
+        "description": tool_description,
+        "name": tool_name,
+        "tool_class": f"{tool_name.capitalize()}Tool"
+    }
+    tool_config_json = json.dumps(tool_config, indent=4)
+
+    # Generate app.py content
+    app_py_content = f'''from transformers.tools.base import launch_gradio_demo
+                    from {tool_name} import {tool_name.capitalize()}Tool
+                    
+                    launch_gradio_demo({tool_name.capitalize()}Tool)
+                    '''
+
+    # Generate requirements.txt content
+    requirements_content = '''transformers>=4.29.0
+            # diffusers
+            accelerate
+            torch
+            '''
+
+    # Generate text_generator.py content
+    text_generator_py_content = f'''import os
+    from transformers import pipeline
+    from transformers import Tool
+    
+    class {tool_name.capitalize()}Tool(Tool):
+        name = "{tool_name}"
+        description = (
+            "{tool_description}"
+        )
+    
+        inputs = ["text"]
+        outputs = ["text"]
+    
+        def __call__(self, prompt: str):
+            token = os.environ['hf']
+            text_generator = pipeline(model="microsoft/Orca-2-13b", token=token)
+            generated_text = text_generator(prompt, max_length=500, num_return_sequences=1, temperature=0.7)
+            print(generated_text)
+            return generated_text
+    '''
+
+    # Write content to files
+    with open("README.md", "w") as readme_file:
+        readme_file.write(readme_content)
+
+    with open("tool_config.json", "w") as tool_config_file:
+        tool_config_file.write(tool_config_json)
+
+    with open("app.py", "w") as app_py_file:
+        app_py_file.write(app_py_content)
+
+    with open("requirements.txt", "w") as requirements_file:
+        requirements_file.write(requirements_content)
+
+    with open(f"{tool_name}.py", "w") as text_generator_py_file:
+        text_generator_py_file.write(text_generator_py_content)
+
+    # Return the generated files for download
+    return "README.md", "tool_config.json", "app.py", "requirements.txt", f"{tool_name}.py"
+
+
+# Define the inputs for the Gradio interface
+io = gr.Interface(generate_files, 
+                  inputs=["text", "text", "text", "text", "text", "text", "text", "text", "checkbox", "text", "text"],
+                  outputs=["text", "text", "text", "text", "text"])
+
+# Launch the Gradio interface
+io.launch()