Supra-Mini-v3-0.5M / inference.py
LH-Tech-AI's picture
Upload 3 files
aedb965 verified
raw
history blame contribute delete
939 Bytes
print("[*] Loading libraries...")
import torch
from transformers import LlamaForCausalLM, PreTrainedTokenizerFast
model_path = "./Supra-Mini-v3-0.5M-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 "
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=150,
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))