coffeeee commited on
Commit
fe765e3
1 Parent(s): 5b26870

switch to using hf model

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -14,8 +14,8 @@ sentence_detector = nltk.data.load('tokenizers/punkt/english.pickle')
14
  tokenizer = GPT2Tokenizer.from_pretrained("gpt2-medium")
15
  tokenizer.truncation_side = 'right'
16
 
17
- model = GPT2LMHeadModel.from_pretrained('checkpoint-10000')
18
- # model = GPT2LMHeadModel.from_pretrained('coffeeee/nsfw-story-generator')
19
  generation_config = GenerationConfig.from_pretrained('gpt2-medium')
20
  generation_config.max_new_tokens = response_length
21
  generation_config.pad_token_id = generation_config.eos_token_id
 
14
  tokenizer = GPT2Tokenizer.from_pretrained("gpt2-medium")
15
  tokenizer.truncation_side = 'right'
16
 
17
+ # model = GPT2LMHeadModel.from_pretrained('checkpoint-10000')
18
+ model = GPT2LMHeadModel.from_pretrained('coffeeee/nsfw-story-generator')
19
  generation_config = GenerationConfig.from_pretrained('gpt2-medium')
20
  generation_config.max_new_tokens = response_length
21
  generation_config.pad_token_id = generation_config.eos_token_id