Spaces:
Sleeping
Sleeping
import base64 | |
import streamlit as st | |
import torch | |
from transformers import GPT2LMHeadModel, GPT2Tokenizer | |
from model.funcs import execution_time | |
def get_base64(file_path): | |
with open(file_path, "rb") as file: | |
base64_bytes = base64.b64encode(file.read()) | |
base64_string = base64_bytes.decode("utf-8") | |
return base64_string | |
def set_background(png_file): | |
bin_str = get_base64(png_file) | |
page_bg_img = ( | |
""" | |
<style> | |
.stApp { | |
background-image: url("data:image/png;base64,%s"); | |
background-size: cover; | |
} | |
</style> | |
""" | |
% bin_str | |
) | |
st.markdown(page_bg_img, unsafe_allow_html=True) | |
set_background("text_generation.png") | |
def load_model(): | |
model_path = "17/" | |
model_name = "sberbank-ai/rugpt3small_based_on_gpt2" | |
tokenizer = GPT2Tokenizer.from_pretrained(model_name) | |
model = GPT2LMHeadModel.from_pretrained(model_path) | |
return tokenizer, model | |
tokenizer, model = load_model() | |
def generate_text( | |
prompt, num_beams=2, temperature=1.5, top_p=0.9, top_k=3, max_length=150 | |
): | |
prompt = tokenizer.encode(prompt, return_tensors="pt") | |
model.eval() | |
with torch.no_grad(): | |
out = model.generate( | |
prompt, | |
do_sample=True, | |
num_beams=num_beams, | |
temperature=temperature, | |
top_p=top_p, | |
top_k=top_k, | |
max_length=max_length, | |
) | |
out = list(map(tokenizer.decode, out))[0] | |
return out | |
with st.sidebar: | |
num_beams = st.slider("Number of Beams", min_value=1, max_value=5, value=2) | |
temperature = st.slider("Temperature", min_value=0.1, max_value=2.0, value=1.5) | |
top_p = st.slider("Top-p", min_value=0.1, max_value=1.0, value=0.9) | |
top_k = st.slider("Top-k", min_value=1, max_value=10, value=3) | |
max_length = st.slider("Maximum Length", min_value=20, max_value=300, value=150) | |
styled_text = """ | |
<style> | |
.styled-text { | |
font-size: 30px; | |
text-shadow: -2px -2px 4px #000000; | |
color: #FFFFFF; | |
-webkit-text-stroke-width: 1px; | |
-webkit-text-stroke-color: #000000; | |
} | |
</style> | |
""" | |
st.markdown(styled_text, unsafe_allow_html=True) | |
prompt = st.text_input( | |
"Ask a question", | |
key="question_input", | |
placeholder="Type here...", | |
type="default", | |
value="", | |
) | |
generate = st.button("Generate", key="generate_button") | |
if generate: | |
if not prompt: | |
st.write("42") | |
else: | |
generated_text = generate_text( | |
prompt, num_beams, temperature, top_p, top_k, max_length | |
) | |
paragraphs = generated_text.split("\n") | |
styled_paragraphs = [ | |
f'<div class="styled-text">{paragraph}</div>' for paragraph in paragraphs | |
] | |
styled_generated_text = " ".join(styled_paragraphs) | |
st.markdown(styled_generated_text, unsafe_allow_html=True) | |