blob: 3290116880f8baf7f6b3903494c92f44738586f9 (
plain) (
blame)
1
2
3
4
5
6
7
8
9
10
|
#!/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)
|
