text-generation-tool / text_generator.py
Chris4K's picture
Update text_generator.py
dfb0190 verified
raw
history blame
3.19 kB
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()