GPT2LMHeadModel / inference.py
Ananthusajeev190's picture
Upload 10 files
2d7594e verified
from transformers import GPT2LMHeadModel, GPT2TokenizerFast
MODEL_DIR = "./68h"
tokenizer = GPT2TokenizerFast.from_pretrained(MODEL_DIR)
model = GPT2LMHeadModel.from_pretrained(MODEL_DIR)
while True:
prompt = input(">>> ")
if prompt == "exit":
break
inputs = tokenizer(prompt, return_tensors="pt")
outputs = model.generate(
**inputs,
max_new_tokens=150,
do_sample=True,
temperature=0.8
)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))