Update app.py

app.py

@@ -6,6 +6,7 @@ import torch
 
 model_checkpoint = "vives/distilbert-base-uncased-finetuned-imdb-accelerate"
 model = AutoModelForMaskedLM.from_pretrained(model_checkpoint,output_hidden_states=True)
+model_base = AutoModelForMaskedLM.from_pretrained("distilbert-base-uncased")
 tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
 text1 = st.text_area("Enter first sentence")
 text2 = st.text_area("Enter second sentence")
@@ -37,10 +38,21 @@ def pool_embeddings(out, tok):
   return mean_pooled
   
 if text1 and text2:
-  tokens = concat_tokens(text1,text2)
-  outputs = model(**tokens)
-  mean_pooled = pool_embeddings(outputs,tokens).detach().numpy()
-  st.write(cosine_similarity(
-      [mean_pooled[0]],
-      mean_pooled[1:]
-  ))
+  with torch.no_grad():
+    tokens = concat_tokens(text1,text2)
+    outputs = model(**tokens)
+    mean_pooled = pool_embeddings(outputs,tokens).detach().numpy()
+    fine_tuned_out = cosine_similarity(
+        [mean_pooled[0]],
+        mean_pooled[1:]
+    )[0]
+    
+    outputs_base = model_base(**tokens)
+    mean_pooled_base = pool_embeddings(outputs_base,tokens).detach().numpy()
+    base_out = cosine_similarity(
+        [mean_pooled[0]],
+        mean_pooled[1:]
+    )[0]
+    
+    st.write(f">>>Similarity for fine-tuned {fine_tuned_out}")
+    st.write(f">>>Similarity for base {base_out}")