File size: 2,866 Bytes
3a905e4 eb91edf c747562 eb91edf 3a905e4 cb2adb5 eb91edf cb2adb5 c747562 3a905e4 eb91edf 3a905e4 eb91edf 3a905e4 eb91edf c747562 3a905e4 c747562 3a905e4 c747562 3a905e4 |
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 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 |
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")
@st.cache_data
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()
@execution_time
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)
|