Update main.py

main.py

@@ -24,12 +24,11 @@ logging.basicConfig(level=logging.INFO)
 ####################################
 # Text Generation Endpoint
 ####################################
-import tempfile
-os.environ["TRANSFORMERS_CACHE"] = tempfile.mkdtemp()
+
 
 HF_TOKEN = os.getenv("manhal")
 
-TEXT_MODEL_NAME = "aubmindlab/aragpt2-base"
+TEXT_MODEL_NAME = "aubmindlab/aragpt2-mega"
 text_tokenizer = AutoTokenizer.from_pretrained(TEXT_MODEL_NAME, token = HF_TOKEN)
 text_model = AutoModelForCausalLM.from_pretrained(TEXT_MODEL_NAME,  token = HF_TOKEN)