File size: 1,961 Bytes
fd3090c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | """
Inference: load the trained decoder and generate from a prompt.
Run: python generate.py "your problem text here"
python generate.py # runs the built-in examples
"""
import sys
import torch
from data import MAX_SRC_LEN
from encoder_loader import load_encoder
from model import SemanticConditionedDecoder
DECODER_WEIGHTS = "decoder_weights.pt"
def load_model(device="cuda", weights=DECODER_WEIGHTS):
encoder, tokenizer = load_encoder(device=device)
model = SemanticConditionedDecoder(encoder=encoder, tokenizer=tokenizer).to(device)
state = torch.load(weights, map_location=device)
model.load_state_dict(state, strict=False)
model.eval()
return model, tokenizer
@torch.no_grad()
def generate_text(model, tokenizer, prompt, device="cuda",
max_new_tokens=256, temperature=0.8):
enc = tokenizer(prompt, padding=True, truncation=True,
max_length=MAX_SRC_LEN, return_tensors="pt")
input_ids = enc["input_ids"].to(device)
attention_mask = enc["attention_mask"].to(device)
out = model.generate(input_ids, attention_mask,
max_new_tokens=max_new_tokens, temperature=temperature)
return tokenizer.decode(out[0][1:], skip_special_tokens=True) # drop BOS
EXAMPLES = [
"252 fifth-grade students and 8 teachers are going on a field trip. "
"If renting a 41-seater bus costs 300,000 won and the highway toll per bus "
"is 7,500 won, how much does it cost to rent the buses and pay the tolls?",
]
def main():
device = "cuda" if torch.cuda.is_available() else "cpu"
model, tokenizer = load_model(device=device)
prompts = sys.argv[1:] or EXAMPLES
for i, prompt in enumerate(prompts, 1):
print("=" * 80)
print(f"PROMPT {i}: {prompt}")
print("-" * 80)
print(generate_text(model, tokenizer, prompt, device=device))
print("=" * 80)
if __name__ == "__main__":
main()
|