add genre to generation

app.py

@@ -3,7 +3,7 @@ from transformers import AutoTokenizer, TFAutoModelForCausalLM
 
 # MODEL TO CALL
 
-model_name = "Alirani/distilgpt2-finetuned-synopsis"
+model_name = "Alirani/distilgpt2-finetuned-synopsis-genres_final"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = TFAutoModelForCausalLM.from_pretrained(model_name)
 
@@ -11,7 +11,7 @@ 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)
-    processed_synopsis = "".join(synopsis.split(':')[1].rpartition('.')[:2]).strip()
+    processed_synopsis = "".join(synopsis.split('|')[2].rpartition('.')[:2]).strip()
     return processed_synopsis
 
 favicon = "https://i.ibb.co/JRdhFZg/favicon-32x32.png"
@@ -24,11 +24,16 @@ st.header('Generate a story')
 
 prod_title = st.text_input('Type a title to generate a synopsis')
 
+option_genres = st.selectbox(
+    'Select a genre to tailor your synopsis',
+    ('Family', 'Romance', 'Comedy', 'Action', 'Documentary', 'Adventure', 'Drama', 'Mystery', 'Crime', 'Thriller', 'Science Fiction', 'History', 'Music', 'Western', 'Fantasy', 'TV Movie', 'Horror', 'Animation', 'Reality')
+    )
+
 button_synopsis = st.button('Get synopsis')
 
 if button_synopsis:
     if len(prod_title.split(' ')) > 0:
-        gen_synopsis = generate_synopsis(model, tokenizer, f"{prod_title} : ")
+        gen_synopsis = generate_synopsis(model, tokenizer, f"{prod_title} | {option_genres} | ")
         st.text_area('Generated synopsis', value=gen_synopsis, disabled=True)
     else:
         st.write('Write a title for the generator to work !')