ChihChiu29 commited on
Commit
8569dd4
1 Parent(s): b0bfe18
Files changed (1) hide show
  1. main.py +4 -4
main.py CHANGED
@@ -14,8 +14,8 @@ token_size_limit = None
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).
@@ -23,8 +23,8 @@ token_size_limit = None
23
  # model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-small")
24
  # tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-base")
25
  # model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-base")
26
- tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-large")
27
- model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-large")
28
  token_size_limit = 512
29
 
30
  # Too large for 16GB
 
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).
 
23
  # model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-small")
24
  # tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-base")
25
  # model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-base")
26
+ # tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-large")
27
+ # model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-large")
28
  token_size_limit = 512
29
 
30
  # Too large for 16GB