Spaces:
Sleeping
Sleeping
File size: 1,351 Bytes
f0350df 53ec071 f0350df 5818a55 f0350df d4379ad f0350df d4379ad f0350df cec4192 f0350df |
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 |
import gradio as gr
import torch
import numpy as np
from model import Transformer
from transformers import AutoTokenizer # pip install transformers
from utils import (
DEVICE,
DROPOUT,
NUM_EMBED,
NUM_HEAD,
NUM_LAYER,
BLOCK_SIZE,
encode,
decode
)
tokenizer = AutoTokenizer.from_pretrained("Neu256/PromeTokenizer")
vocab_size = tokenizer.vocab_size
# train a new model
model = Transformer(
vocab_size=vocab_size,
num_embed=NUM_EMBED,
block_size=BLOCK_SIZE,
num_heads=NUM_HEAD,
num_layers=NUM_LAYER,
dropout=DROPOUT
)
# load model to GPU if available
m = model.to(DEVICE)
# print the number of parameters in the model
m.load_state_dict(torch.load("base_model_1.pth", map_location=torch.device(DEVICE)))
m.eval()
#print(
# "Model with {:.2f}M parameters".format(sum(p.numel() for p in m.parameters()) / 1e6)
#)
def model_generate(text, number_of_new_token, temperature, top_p):
context = encode(str(text), tokenizer).unsqueeze(0).to(DEVICE)
return decode(enc_sec=m.generate(idx=context, max_new_tokens=number_of_new_token, temperature = temperature, top_p = top_p)[0], tokenizer=tokenizer)
iface = gr.Interface(fn=model_generate, inputs=["text", gr.Slider(10, 1000), gr.Slider(0, 1, value=0.7, step = 0.05), gr.Slider(0, 1, value=0.95, step = 0.05)], outputs="text")
iface.launch() |