add classifier

app.py

@@ -1,11 +1,15 @@
 import streamlit as st
-from transformers import AutoTokenizer, TFAutoModelForCausalLM
+from transformers import AutoTokenizer, TFAutoModelForCausalLM, TFAutoModelForSequenceClassification
 
 # MODEL TO CALL
 
-model_name = "Alirani/distilgpt2-finetuned-synopsis-genres_final"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = TFAutoModelForCausalLM.from_pretrained(model_name)
+generator_name = "Alirani/distilgpt2-finetuned-synopsis-genres_final"
+tokenizer_gen = AutoTokenizer.from_pretrained(generator_name)
+model_gen = TFAutoModelForCausalLM.from_pretrained(generator_name)
+
+classifier_name = "Alirani/overview_classifier_final"
+tokenizer_clf = AutoTokenizer.from_pretrained(classifier_name)
+model_clf = TFAutoModelForSequenceClassification.from_pretrained(classifier_name)
 
 def generate_synopsis(model, tokenizer, title):
     input_ids = tokenizer(title, return_tensors="tf")
@@ -14,6 +18,11 @@ def generate_synopsis(model, tokenizer, title):
     processed_synopsis = "".join(synopsis.split('|')[2].rpartition('.')[:2]).strip()
     return processed_synopsis
 
+def generate_classification(model, tokenizer, title, overview):
+    tokens = tokenizer(f"{title} | {overview}", padding=True, truncation=True, return_tensors="tf")
+    output = model(**tokens)
+    return output
+
 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')
@@ -23,6 +32,7 @@ st.title('Demo LoreFinder')
 st.header('Generate a story')
 
 prod_title = st.text_input('Type a title to generate a synopsis')
+prod_synopsis = st.text_input('Type a synopsis to guess the genre')
 
 option_genres = st.selectbox(
     'Select a genre to tailor your synopsis',
@@ -33,9 +43,18 @@ option_genres = st.selectbox(
 
 button_synopsis = st.button('Get synopsis')
 
+button_genres = st.button('Get genres')
+
 if button_synopsis:
     if len(prod_title.split(' ')) > 0:
-        gen_synopsis = generate_synopsis(model, tokenizer, f"{prod_title} | {option_genres} | ")
+        gen_synopsis = generate_synopsis(model_gen, tokenizer_gen, 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 !')
+
+if button_genres:
+    if len(prod_title.split(' ')) > 0:
+        classifier_output = generate_classification(model_clf, tokenizer_clf, prod_title, prod_synopsis)
+        st.write(f"Guessed genre : {classifier_output}")
+    else:
+        st.write('Write a title and a synopsis for the classifier to work !')