Views
No views yet
model.py file in your directory (which is included in this repo).1import torch
2import json
3from tokenizers import ByteLevelBPETokenizer
4from model import SLM, SLMConfig # Requires model.py in the same folder
5
6# 1. Load Configuration
7with open("config.json", "r") as f:
8 config_dict = json.load(f)
9config = SLMConfig(**config_dict)
10
11# 2. Initialize Model Architecture & Load Weights
12model = SLM(config)
13# Ensure you match the filename you uploaded (vital_lm_25m_swiglu_best.pt)
14model.load_state_dict(torch.load("vital_lm_25m_swiglu_best.pt", map_location='cpu'))
15model.eval()
16
17# 3. Load Tokenizer (Crucial Step!)
18# The tokenizer files (vocab.json, merges.txt) are in this repo
19tokenizer = ByteLevelBPETokenizer(
20 "updated_vocab.json",
21 "updated_merges.txt"
22)
23
24# 4. Chat Function
25def chat(text, max_new_tokens=50):
26 # Encode Input
27 ids = tokenizer.encode(text).ids
28 idx = torch.tensor(ids).unsqueeze(0)
29
30 # Generate
31 with torch.no_grad():
32 out = model.generate(idx, max_new_tokens=max_new_tokens, temperature=0.3, top_k=40)
33
34 # Decode Output
35 print(tokenizer.decode(out[0].tolist()))
36
37# Example Test
38chat("Patient: I have a severe headache and sensitivity to light. Doctor:")