ChihChiu29 commited on
Commit
23d399e
·
1 Parent(s): 2f79a00

switch to 400M model

Browse files
Files changed (1) hide show
  1. main.py +5 -4
main.py CHANGED
@@ -12,10 +12,11 @@ from transformers import T5Tokenizer, T5ForConditionalGeneration
12
  token_size_limit = None
13
 
14
  # FROM: https://huggingface.co/facebook/blenderbot-400M-distill?text=Hey+my+name+is+Thomas%21+How+are+you%3F
15
- # tokenizer = AutoTokenizer.from_pretrained("facebook/blenderbot-400M-distill")
16
- # model = AutoModelForSeq2SeqLM.from_pretrained("facebook/blenderbot-400M-distill")
17
- tokenizer = AutoTokenizer.from_pretrained("facebook/blenderbot-1B-distill")
18
- model = AutoModelForSeq2SeqLM.from_pretrained("facebook/blenderbot-1B-distill")
 
19
  token_size_limit = 128
20
 
21
  # T5 model can use "any" sequence lenghth, but memory usage is O(L^2).
 
12
  token_size_limit = None
13
 
14
  # FROM: https://huggingface.co/facebook/blenderbot-400M-distill?text=Hey+my+name+is+Thomas%21+How+are+you%3F
15
+ tokenizer = AutoTokenizer.from_pretrained("facebook/blenderbot-400M-distill")
16
+ model = AutoModelForSeq2SeqLM.from_pretrained(
17
+ "facebook/blenderbot-400M-distill")
18
+ # tokenizer = AutoTokenizer.from_pretrained("facebook/blenderbot-1B-distill")
19
+ # model = AutoModelForSeq2SeqLM.from_pretrained("facebook/blenderbot-1B-distill")
20
  token_size_limit = 128
21
 
22
  # T5 model can use "any" sequence lenghth, but memory usage is O(L^2).