Supra-Mini-v5-8M / inference.py
LH-Tech-AI's picture
Upload 11 files
d467b07 verified
raw
history blame
1.03 kB
print("[*] Loading libraries...")
import torch
from transformers import LlamaForCausalLM, PreTrainedTokenizerFast
model_path = "./Supra-Mini-v5-8M-FINAL"
print("[*] Loading tokenizer...")
tokenizer = PreTrainedTokenizerFast.from_pretrained(model_path)
print("[*] Loading model...")
model = LlamaForCausalLM.from_pretrained(model_path)
model.eval()
prompt = "The main concept of physics is " # "Artificial intelligence is " | "The main concept of physics is " | "Once upon a time, "
print(f"[*] Prompt: {prompt!r}")
inputs = tokenizer(prompt, return_tensors="pt")
with torch.no_grad():
outputs = model.generate(
input_ids=inputs["input_ids"],
attention_mask=inputs["attention_mask"],
max_new_tokens=256,
do_sample=True,
temperature=0.5,
top_p=0.9,
top_k=25,
repetition_penalty=1.3,
pad_token_id=tokenizer.pad_token_id,
eos_token_id=tokenizer.eos_token_id,
)
print("[*] Output:", tokenizer.decode(outputs[0], skip_special_tokens=True))