Spaces:
Runtime error
Runtime error
File size: 4,127 Bytes
7b715b6 99f39f4 7b715b6 99f39f4 7b715b6 99f39f4 7b715b6 46e6863 7b715b6 99f39f4 7b715b6 99f39f4 7b715b6 99f39f4 7b715b6 4da6cf6 99f39f4 | 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 | import time
import torch
import torch.nn.functional as F
import streamlit as st
from parser import extract_titles
from generator import (
get_model_and_tokenizer,
generate_prompt,
generate_comment,
extract_comment_from_prompt
)
@st.cache_resource
def preprocess(model_name, model_type="gpt2"):
tokenizer, model = get_model_and_tokenizer("cpu", "./chkp", model_name, model_type=model_type)
return tokenizer, model
# Greetings
st.title("Generate YouTube Comments")
st.write("This mini-app generates English comments for a YouTube video. It uses a fine-tuned version of GPT-2 by OpenAI. You can find the code in Files section.")
# Pre-process
tokenizer, model = preprocess("GPT2_02_Ep0_St300000.pt", model_type="gpt2")
model.eval()
filter_value = -float("Inf")
entry_length = 20
# Settings
input_url = st.text_input(label="Drop URL below", placeholder="https://youtu.be/mCV44C5rQ2M")
col1, col2 = st.columns(2)
with col1:
temperature = st.slider('Temperature of sampling', 0.0, 1.0, value=0.7, key=7)
with col2:
top_p = st.slider('Top-p parameter', 0.0, 1.0, value=0.8, key=8)
col1, _ = st.columns(2)
with col1:
num_of_coms = st.slider('Number of comments', 1, 5, value=1, key=9)
# Action
if st.button("Generate text", type="primary"):
if len(input_url):
try:
# generate & print
channel, title = extract_titles(input_url)
prompt = generate_prompt(title)
with torch.no_grad():
raw_generated = torch.tensor(tokenizer.encode(prompt)).unsqueeze(0)
st.markdown("""---""")
place_holders = []
texts = []
for index in range(num_of_coms):
with st.spinner('Please wait while your comment is being generated...'):
place_holders.append(st.empty())
generated = raw_generated.clone()
with torch.no_grad():
for i in range(entry_length):
outputs = model(generated, labels=generated)
loss, logits = outputs[:2]
logits = logits[:, -1, :] / (temperature if temperature > 0 else 1.0)
sorted_logits, sorted_indices = torch.sort(logits, descending=True)
cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
sorted_indices_to_remove = cumulative_probs > top_p
sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
sorted_indices_to_remove[..., 0] = 0
indices_to_remove = sorted_indices[sorted_indices_to_remove]
logits[:, indices_to_remove] = filter_value
next_token = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)
generated = torch.cat((generated, next_token), dim=1)
output_list = list(generated.squeeze().numpy())
output_text = f"{tokenizer.decode(output_list)} <EOS>"
output_text = extract_comment_from_prompt(output_text)
place_holders[index].text_area(label=f"Comment #{index + 1}",
value=output_text, disabled=True, key=(index + 1)*1000+i)
if next_token in tokenizer.encode("<EOS>"):
break
output_list = list(generated.squeeze().numpy())
output_text = f"{tokenizer.decode(output_list)} <EOS>"
texts.append(extract_comment_from_prompt(output_text))
for index, text in enumerate(texts):
place_holders[index].text_area(label=f"Comment #{index + 1}", value=text, disabled=False, key=index+100)
except RuntimeError as e:
st.error(e.args[0], icon="🚨")
else:
st.error('Please enter a url', icon="🚨") |