free-x commited on
Commit
a2ec2a8
1 Parent(s): 1213d0b

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +43 -0
app.py ADDED
@@ -0,0 +1,43 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import subprocess
3
+ import logging
4
+
5
+ import gradio as gr
6
+ from transformers import pipeline
7
+ from transformers import AutoTokenizer, AutoModelForCausalLM
8
+
9
+ # 下载模型
10
+ base_dir = "/root/.cache/huggingface/hub"
11
+ if not os.path.isdir(base_dir):
12
+ os.makedirs(base_dir)
13
+
14
+ cmd_list = ["cd", base_dir, "&&", "git lfs install", "&&", "git clone", "https://gitee.com/lanzhiwang/gpt2.git", "models"]
15
+ cmd_str = " ".join(cmd_list)
16
+ print("cmd_str:", cmd_str)
17
+ ret, out = subprocess.getstatusoutput(cmd_str)
18
+ print("ret:", ret)
19
+ print("out:", out)
20
+
21
+ tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path="/root/.cache/huggingface/hub/models")
22
+ model = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path="/root/.cache/huggingface/hub/models")
23
+ generator = pipeline('text-generation', model=model, tokenizer=tokenizer)
24
+
25
+ # generator = pipeline('text-generation', model='gpt2')
26
+
27
+ def generate(text):
28
+ result = generator(text, max_length=30, num_return_sequences=1)
29
+ return result[0]["generated_text"]
30
+
31
+ examples = [
32
+ ["The Moon's orbit around Earth has"],
33
+ ["The smooth Borealis basin in the Northern Hemisphere covers 40%"],
34
+ ]
35
+
36
+ demo = gr.Interface(
37
+ fn=generate,
38
+ inputs=gr.inputs.Textbox(lines=5, label="Input Text"),
39
+ outputs=gr.outputs.Textbox(label="Generated Text"),
40
+ examples=examples
41
+ )
42
+
43
+ demo.launch(server_name="0.0.0.0", server_port=7860)