When trying to learn AI I realized that loading the model on a smaller and less capable device would be necessary. Considering I would like t o hypothetically load the model up on my phone or a cheep laptop. I first started googling around, but that got me no where considering i did not have the structure set up to save the file. First you need a model in a nn.Module class. If you have a bunch of nn.Layers you could possibly load them all individually, but that might not be desired, instead you can connect them ll together with the nn.Sequential command. This strings the output of one class into the input of the next layer and returns the result of the last layer.
Code
1import torch
2from torch import tensor
3
4import torch.nn.functional as F
5import torch.nn as nn
6from transformers import AutoTokenizer
7
8# The model needed a tokenizer, so Itrained it on this one
9tokenizer = AutoTokenizer.from_pretrained("gpt2", add_prefix_space=True)
10
11# This selects the CPU to load the data to
12device = torch.device('cpu')
13
14# These lines are just the initilization lines for the model I created.
15inner = 700
16
17model = nn.Sequential(
18 nn.Linear(size,inner, bias=False),
19 nn.Linear(inner,size, bias=False)
20)
21
22# This actually loades the model from the file.
23model.load_state_dict(torch.load("model.pt", map_location=device))
24
25# This is just a testing loop.
26while True:
27 string = input("Enter some text: ")
28 one_hot = F.one_hot(tensor(tokenizer.encode(string)),size).float()
29 results = model(one_hot)
30 output = results.argmax(dim=1)
31 term = tokenizer.decode(output)
32 print(term)