lanzhiwang commited on
Commit
3c3d1c4
1 Parent(s): 0a459ec

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +10 -13
app.py CHANGED
@@ -7,19 +7,16 @@ 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 clone", "https://gitee.com/lanzhiwang/gpt2.git", "models"]
15
- cmd_str = "".join(cmd_list)
16
- logging.info("cmd_str: %s", cmd_str)
17
- ret, out = subprocess.getstatusoutput(cmd_str)
18
- logging.info("ret: %s", ret)
19
- logging.info("out: %s", 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')
 
7
  from transformers import AutoTokenizer, AutoModelForCausalLM
8
 
9
  # 下载模型
10
+
11
+
12
+
13
+ # Load model directly
14
+ from transformers import AutoTokenizer, AutoModelForCausalLM
15
+
16
+ tokenizer = AutoTokenizer.from_pretrained("gpt2")
17
+ model = AutoModelForCausalLM.from_pretrained("gpt2")
18
+
19
+
 
 
 
20
  generator = pipeline('text-generation', model=model, tokenizer=tokenizer)
21
 
22
  # generator = pipeline('text-generation', model='gpt2')