1pip install transformers torch accelerate
2
3Method 1: Direct Transformers Loading
4from transformers import AutoTokenizer, AutoModelForCausalLM
5import torch
6
7# Load model and tokenizer
8model_id = "arif-butt/tinyllama-trl-merged"
9tokenizer = AutoTokenizer.from_pretrained(model_id)
10model = AutoModelForCausalLM.from_pretrained(
11 model_id,
12 torch_dtype=torch.float16,
13 device_map="auto",
14 trust_remote_code=True,
15)
16model.eval()
17
18# Define prompt template
19prompt = "Q: What is machine learning?\nA:"
20
21# Tokenize
22inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
23
24# Generate
25with torch.no_grad():
26 outputs = model.generate(
27 **inputs,
28 max_new_tokens=150,
29 temperature=0.7,
30 top_p=0.95,
31 do_sample=True,
32 repetition_penalty=1.1,
33 pad_token_id=tokenizer.eos_token_id,
34 )
35
36# Decode and print
37response = tokenizer.decode(outputs[0], skip_special_tokens=True)
38print(f"Prompt: {prompt}")
39print(f"Response: {response[len(prompt):].strip()}")
40
41Method 2: Pipeline for Simple Inference
42from transformers import pipeline
43
44pipe = pipeline(
45 "text-generation",
46 model="arif-butt/tinyllama-trl-merged",
47 torch_dtype=torch.float16,
48 device_map="auto",
49)
50
51prompt = "Q: Explain neural networks in simple terms\nA:"
52result = pipe(prompt, max_new_tokens=150, temperature=0.7, do_sample=True)
53print(result[0]["generated_text"])
54