File size: 1,556 Bytes
d171a71 |
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 |
"""
Example script for running inference with the Hugging Face model.
"""
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import warnings
# Suppress the TypedStorage deprecation warning
warnings.filterwarnings('ignore', category=UserWarning, message='.*TypedStorage is deprecated.*')
def main():
model_path = "." # Path to the model, or use "YOUR_USERNAME/YOUR_MODEL_NAME" for HF Hub
print("Loading model and tokenizer...")
model = AutoModelForCausalLM.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
model.eval()
print(f"Model loaded on {device}")
print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}")
prompts = [
"Once upon a time",
"The quick brown fox",
]
for prompt in prompts:
print(f"\n{'='*60}")
print(f"Prompt: {prompt}")
print(f"{'='*60}")
inputs = tokenizer(prompt, return_tensors="pt").to(device)
with torch.no_grad():
outputs = model.generate(
**inputs,
max_length=100,
temperature=1.0,
top_k=50,
top_p=0.9,
do_sample=True,
pad_token_id=tokenizer.eos_token_id
)
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(f"\nGenerated:\n{generated_text}")
if __name__ == "__main__":
main()
|