import os import subprocess import gradio as gr from transformers import pipeline from transformers import AutoTokenizer, AutoModelForCausalLM # 下载模型 base_dir = "/root/.cache/huggingface/hub" if not os.path.isdir(base_dir): os.makedirs(base_dir) cmd_list = ["cd", base_dir, "&&", "git lfs install", "&&", "git clone", "https://gitee.com/lanzhiwang/gpt2.git", "models"] cmd_str = " ".join(cmd_list) print("cmd_str:", cmd_str) ret, out = subprocess.getstatusoutput(cmd_str) print("ret:", ret) print("out:", out) tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path="/root/.cache/huggingface/hub/models") model = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path="/root/.cache/huggingface/hub/models") generator = pipeline('text-generation', model=model, tokenizer=tokenizer) # generator = pipeline('text-generation', model='gpt2') def generate(text): result = generator(text, max_length=30, num_return_sequences=1) return result[0]["generated_text"] examples = [ ["The Moon's orbit around Earth has"], ["The smooth Borealis basin in the Northern Hemisphere covers 40%"], ] demo = gr.Interface( fn=generate, inputs=gr.inputs.Textbox(lines=5, label="Input Text"), outputs=gr.outputs.Textbox(label="Generated Text"), examples=examples ) demo.launch(server_name="0.0.0.0", server_port=7860)