|
|
import streamlit as st |
|
|
from transformers import AutoTokenizer, TFAutoModelForCausalLM |
|
|
|
|
|
|
|
|
|
|
|
model_name = "Alirani/distilgpt2-finetuned-synopsis" |
|
|
tokenizer = AutoTokenizer.from_pretrained(model_name) |
|
|
model = TFAutoModelForCausalLM.from_pretrained(model_name) |
|
|
|
|
|
def generate_synopsis(model, tokenizer, title): |
|
|
input_ids = tokenizer(title, return_tensors="tf") |
|
|
output = model.generate(input_ids['input_ids'], max_length=150, num_beams=5, no_repeat_ngram_size=2, top_k=50, attention_mask=input_ids['attention_mask']) |
|
|
synopsis = tokenizer.decode(output[0], skip_special_tokens=True) |
|
|
return synopsis |
|
|
|
|
|
favicon = "https://i.ibb.co/JRdhFZg/favicon-32x32.png" |
|
|
|
|
|
st.set_page_config(page_title="LoreFinder-demo", page_icon = favicon, layout = 'wide', initial_sidebar_state = 'auto') |
|
|
|
|
|
st.title('Demo LoreFinder') |
|
|
|
|
|
st.header('Generate a story') |
|
|
|
|
|
prod_title = st.text_input('Type a title to generate a synopsis') |
|
|
|
|
|
button_synopsis = st.button('Get synopsis') |
|
|
|
|
|
if button_synopsis: |
|
|
if len(prod_title.split(' ')) > 0: |
|
|
gen_synopsis = generate_synopsis(model, tokenizer, prod_title) |
|
|
st.text_area(gen_synopsis, disabled=True) |
|
|
else: |
|
|
st.write('Write a title for the generator to work !') |
|
|
|
