#!/usr/bin/python3
from aitextgen import aitextgen
import sys

#ai = aitextgen(model_folder="/home/msglm/Documents/ComputerScience/Projects/GPT-2/RedpillGenerator/aitextgen/")
#ai = aitextgen(model="distilgpt2", cache_dir="/home/msglm/.local/share/aitextgen-models")
ai = aitextgen(tf_gpt2="774M", cache_dir="/home/msglm/.local/share/aitextgen-models")
ai.to_gpu()
output = ai.generate_one(max_length=int(sys.argv[1]), prompt=sys.argv[2], temperature=0.7)
print(output)