Update main.py

main.py

@@ -4,6 +4,7 @@ import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForSeq2SeqLM
 import logging
 import re
+import os
 
 app = FastAPI()
 
@@ -23,8 +24,8 @@ logging.basicConfig(level=logging.INFO)
 ####################################
 # Text Generation Endpoint
 ####################################
-
-TEXT_MODEL_NAME = "aubmindlab/aragpt2-medium"
+os.environ["TRANSFORMERS_CACHE"] = "/app/cache"
+TEXT_MODEL_NAME = "aubmindlab/aragpt2-base"
 text_tokenizer = AutoTokenizer.from_pretrained(TEXT_MODEL_NAME)
 text_model = AutoModelForCausalLM.from_pretrained(TEXT_MODEL_NAME)