Ehren12 commited on
Commit
a200dd4
1 Parent(s): 5e2686d

switch back to 90M model

Browse files
Files changed (1) hide show
  1. main.py +3 -3
main.py CHANGED
@@ -1,7 +1,7 @@
1
  from dotenv import load_dotenv
2
  import os
3
 
4
- from transformers import BlenderbotTokenizer, BlenderbotForConditionalGeneration
5
  from helpers import in_cache
6
 
7
  from fastapi import FastAPI
@@ -13,8 +13,8 @@ from starlette.middleware.cors import CORSMiddleware
13
 
14
  load_dotenv()
15
 
16
- model = BlenderbotForConditionalGeneration.from_pretrained("facebook/blenderbot-400M-distill", cache_dir="new_cache_dir/")
17
- tokenizer = BlenderbotTokenizer.from_pretrained("facebook/blenderbot-400M-distill", cache_dir="new_cache_dir/")
18
  class UserMSGRequest(BaseModel):
19
  message: str
20
 
 
1
  from dotenv import load_dotenv
2
  import os
3
 
4
+ from transformers import BlenderbotSmallForConditionalGeneration, AutoTokenizer
5
  from helpers import in_cache
6
 
7
  from fastapi import FastAPI
 
13
 
14
  load_dotenv()
15
 
16
+ model = BlenderbotSmallForConditionalGeneration.from_pretrained("facebook/blenderbot_small-90M", cache_dir="new_cache_dir/")
17
+ tokenizer = AutoTokenizer.from_pretrained("facebook/blenderbot_small-90M", cache_dir="new_cache_dir/")
18
  class UserMSGRequest(BaseModel):
19
  message: str
20