Merge branch 'main' of https://huggingface.co/spaces/A-M-S/movie-genre

app.py

@@ -16,7 +16,7 @@ st.caption("Either enter Wiki URL or the Cast info of the movie. Cast will be fe
 wiki_url = st.text_input("Enter Wiki URL of the movie (Needed for fetching the cast information)")
 cast_input = st.text_input("Enter Wiki IDs of the cast (Should be separated by comma)")
 
-model = AutoModelForSequenceClassification.from_pretrained("./checkpoint-36819")
+model = AutoModelForSequenceClassification.from_pretrained("./checkpoint-49092")
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 model.to(device)
 
@@ -61,7 +61,7 @@ if st.button("Predict"):
     # Use Meta Model approach when cast information is available otherwise use DistilBERT model
     if len(cast)!=0:
             # Base Model 1: DistilBERT
-            id2label, label2id, tokenizer, tokenized_plot = utility.tokenize(clean_plot, ["Action","Drama", "Romance", "Comedy", "Thriller"])
+            id2label, label2id, tokenizer, tokenized_plot = utility.tokenize(clean_plot, ["Action","Drama", "Romance", "Comedy", "Thriller","Crime","Horror"])
             input_ids = [np.asarray(tokenized_plot['input_ids'])]
             attention_mask = [np.asarray(tokenized_plot['attention_mask'])]
         
@@ -80,11 +80,17 @@ if st.button("Predict"):
         
             # Concatenating Outputs of base models
             r1 = distilbert_pred[3]
-            r2 = distilbert_pred[1]
-            r3 = distilbert_pred[2]
+            r2 = distilbert_pred[5]
+            r3 = distilbert_pred[1]
+            r4 = distilbert_pred[6]
+            r5 = distilbert_pred[2]
+            r6 = distilbert_pred[4]
             distilbert_pred[1] = r1
             distilbert_pred[2] = r2
             distilbert_pred[3] = r3
+            distilbert_pred[4] = r4
+            distilbert_pred[5] = r5
+            distilbert_pred[6] = r6
             pred1 = distilbert_pred
             pred2 = lr_model_pred
             distilbert_pred = pred1.detach().numpy()
@@ -95,13 +101,13 @@ if st.button("Predict"):
             probs = meta_model.predict_proba([concat_features])
             
             # Preparing Output
-            id2label = {0:"Action",1:"Comedy",2:"Drama",3:"Romance",4:"Thriller"}
+            id2label = {0: "Action",1: "Comedy",2: "Crime",3: "Drama",4: "Horror",5: "Romance",6: "Thriller"}
             i = 0
             for prob in probs[0]:
                 out.append([id2label[i], prob])
                 i += 1
     else:
-         id2label, label2id, tokenizer, tokenized_plot = utility.tokenize(clean_plot, ["Action","Drama", "Romance", "Comedy", "Thriller"])
+         id2label, label2id, tokenizer, tokenized_plot = utility.tokenize(clean_plot, ["Action","Drama", "Romance", "Comedy", "Thriller","Crime","Horror"])
          input_ids = [np.asarray(tokenized_plot['input_ids'])]
          attention_mask = [np.asarray(tokenized_plot['attention_mask'])]