Views
No views yet
1import mini_models # Register custom Mini-LLM models
2from transformers import AutoTokenizer, AutoModelForCausalLM
3import torch
4
5model = AutoModelForCausalLM.from_pretrained("WKQ9411/Mini-Llama3-100M-Base")
6tokenizer = AutoTokenizer.from_pretrained("WKQ9411/Mini-Llama3-100M-Base")
7device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
8model = model.to(device)
9
10# Generate text
11input_text = "长城是"
12input_ids = tokenizer(input_text, return_tensors="pt")["input_ids"].to(model.device)
13response = model.generate(input_ids, max_new_tokens=100)
14response = tokenizer.decode(response[0][len(input_ids[0]):], skip_special_tokens=True)
15print(response)1from mini_models import get_model_and_config
2from transformers import AutoTokenizer
3import torch
4
5Model, Config = get_model_and_config("mini_llama3")
6model = Model.from_pretrained("path/to/your/model")
7tokenizer = AutoTokenizer.from_pretrained("path/to/your/tokenizer")
8device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
9model = model.to(device)
10
11# Use the model for generation
12input_text = "长城是"
13input_ids = tokenizer(input_text, return_tensors="pt")["input_ids"].to(model.device)
14response = model.generate(input_ids, max_new_tokens=100)
15response = tokenizer.decode(response[0][len(input_ids[0]):], skip_special_tokens=True)
16print(response)