File size: 2,561 Bytes
f7d62af 3ed9d82 f7d62af |
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 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 |
import streamlit as st
import torch
from time import perf_counter
from transformers import AutoTokenizer, AutoModelForCausalLM
MODEL = 'nan-dre/maneleGPT-medium'
TOKENIZER = 'nan-dre/maneleGPT-medium'
MAX_LENGTH = 256
st.set_page_config(
page_title="ManeleGPT",
page_icon="🇷🇴",
layout="centered"
)
def typical_sampling(model, input_ids, attention_mask, no_repeat_ngram_size, max_length, temperature, typical_p):
return model.generate(
input_ids=input_ids,
attention_mask=attention_mask,
no_repeat_ngram_size=no_repeat_ngram_size,
max_length=max_length,
do_sample=True,
temperature=temperature,
typical_p=typical_p,
top_k=0
)
@st.cache(allow_output_mutation=True)
def setModel():
model = AutoModelForCausalLM.from_pretrained(MODEL)
tokenizer = AutoTokenizer.from_pretrained(TOKENIZER)
return model, tokenizer
st.header("ManeleGPT")
temperature = st.slider(label="Temperatura", min_value=0.01, max_value=1.0, value=0.5, step=0.01)
input = st.text_input(label="Cu ce vers sa inceapa maneaua?", value="", key="seed")
if input:
model, tokenizer = setModel()
tokenized_text = tokenizer(input, add_special_tokens=False, return_tensors="pt")
if len(tokenized_text.input_ids[0]) + MAX_LENGTH > 512: # need to keep less words
keep_last = 512 - MAX_LENGTH
print(f"keep last: {keep_last}")
input_ids, attention_mask = tokenized_text.input_ids[0][-keep_last:], tokenized_text.attention_mask[0][-keep_last:]
previous_ids = tokenized_text.input_ids[0][:keep_last]
st.warning(f"kept last {keep_last}")
else:
input_ids, attention_mask = tokenized_text.input_ids[0], tokenized_text.attention_mask[0]
previous_ids = None
length = min(512, len(input_ids) + MAX_LENGTH)
timer_mark = perf_counter()
output = typical_sampling(model, input_ids.unsqueeze(dim=0), attention_mask.unsqueeze(dim=0), no_repeat_ngram_size=2, max_length=MAX_LENGTH, temperature=temperature, typical_p=1)
details = f"Text generated in {perf_counter()-timer_mark:.2f}s"
if previous_ids is not None:
print(f"\nConcat prev id: "+tokenizer.decode(previous_ids, skip_special_tokens=True))
print(f"\nWith current decode: " + tokenizer.decode(output[0], skip_special_tokens=True))
new_text = tokenizer.decode(torch.cat([previous_ids, output[0]], dim=-1), skip_special_tokens=True)
else:
new_text = tokenizer.decode(output[0], skip_special_tokens=True)
st.text(new_text) |