File size: 1,602 Bytes
366d265 | 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 | print("Loading...")
import torch
from transformers import LlamaForCausalLM, PreTrainedTokenizerFast
def run_inference():
model_path = "./StorySupra-10M"
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")
tokenizer = PreTrainedTokenizerFast.from_pretrained(model_path)
model = LlamaForCausalLM.from_pretrained(model_path)
model.to(device)
model.eval()
def generate_text(prompt, max_new_tokens=100, temperature=0.55, top_k=25, top_p=0.85, repetition_penalty=1.1):
inputs = tokenizer(prompt, return_tensors="pt").to(device)
with torch.no_grad():
output_tokens = model.generate(
**inputs,
max_new_tokens=max_new_tokens,
do_sample=True,
temperature=temperature,
top_k=top_k,
top_p=top_p,
repetition_penalty=repetition_penalty,
pad_token_id=tokenizer.pad_token_id,
eos_token_id=tokenizer.eos_token_id
)
return tokenizer.decode(output_tokens[0], skip_special_tokens=True)
print("-" * 30)
print("StorySupra Story Generator loaded!")
print("Enter a prompt (or type 'exit' to quit):")
while True:
user_prompt = input("\nYour prompt: ")
if user_prompt.lower() in ["exit", "quit", "leave"]:
break
story = generate_text(user_prompt)
print(f"\nGenerated story:\n{story}")
print("-" * 20)
if __name__ == "__main__":
run_inference() |