Update app.py

app.py

@@ -1,43 +1,66 @@
 import streamlit as st
-from transformers import RobertaTokenizer,AutoModelForSequenceClassification
+from transformers import RobertaTokenizer, AutoModelForSequenceClassification
 import torch
 
-state_dict=torch.load("fine_tuned_roberta_comments.bin",map_location=torch.device("cpu"))
-tokenizer=RobertaTokenizer.from_pretrained("roberta-base")
-model =  AutoModelForSequenceClassification.from_pretrained('roberta-base',
-                                                      problem_type="multi_label_classification", 
-                                                      num_labels=3               
-                                                     )
-model.load_state_dict(state_dict)
-device = torch.device("cpu")
-model.to(device)
+# Chargement des modèles
+state_dict_comment = torch.load("fine_tuned_roberta_comment.bin", map_location=torch.device("cpu"))
+state_dict_full = torch.load("fine_tuned_roberta_full.bin", map_location=torch.device("cpu"))
 
+tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
+model_comment = AutoModelForSequenceClassification.from_pretrained('roberta-base', problem_type="multi_label_classification", num_labels=3)
+model_comment.load_state_dict(state_dict_comment)
 
-def main():
-    st.title("Classification de séquence")
+model_full = AutoModelForSequenceClassification.from_pretrained('roberta-base', problem_type="multi_label_classification", num_labels=3)
+model_full.load_state_dict(state_dict_full)
 
-    title = st.text_input("Titre")
-    post = st.text_area("Post")
-    comment = st.text_area("Commentaire")
+device = torch.device("cpu")
+model_comment.to(device)
+model_full.to(device)
 
-    if st.button("Tester"):
-        result = get_predictions(title, post, comment)
-        st.success(result)
 @st.cache_data
-def get_predictions(title, post, commentaire):
-    model.eval()
-    inputs = tokenizer("comment: " + commentaire, return_tensors="pt", padding=True, truncation=True, max_length=512)
+def get_predictions_comment(commentaire):
+    model_comment.eval()
+    inputs = tokenizer(commentaire, return_tensors="pt", padding=True, truncation=True, max_length=512)
     input_ids = inputs['input_ids'].to(device)
     attention_mask = inputs['attention_mask'].to(device)
     with torch.no_grad():
-        outputs = model(input_ids, attention_mask=attention_mask)
+        outputs = model_comment(input_ids, attention_mask=attention_mask)
         logits = outputs.logits
+        probs = torch.softmax(logits, dim=1)
         _, preds = torch.max(logits, dim=1)
-    id2label = {
-        0: "neutral",
-        1: "with palestine",
-        2: "with israel"
-    }
-    return id2label[preds.item()]
+    id2label = {0: "neutral", 1: "with palestine", 2: "with israel"}
+    return id2label[preds.item()], probs.squeeze().tolist()
+
+@st.cache_data
+def get_predictions_full(title, post, commentaire):
+    model_full.eval()
+    inputs = tokenizer("title of the post: " + title + "\n" + "post: " + post + "\n" + "comment: " + commentaire, return_tensors="pt", padding=True, truncation=True, max_length=512)
+    input_ids = inputs['input_ids'].to(device)
+    attention_mask = inputs['attention_mask'].to(device)
+    with torch.no_grad():
+        outputs = model_full(input_ids, attention_mask=attention_mask)
+        logits = outputs.logits
+        probs = torch.softmax(logits, dim=1)
+        _, preds = torch.max(logits, dim=1)
+    id2label = {0: "neutral", 1: "with palestine", 2: "with israel"}
+    return id2label[preds.item()], probs.squeeze().tolist()
+
+def main():
+    st.title("Classification de séquence")
+    title = st.text_input("Titre")
+    post = st.text_area("Post")
+    comment = st.text_area("Commentaire")
+    if st.button("Tester"):
+        if title or post:
+            result, probs = get_predictions_full(title, post, comment)
+        else:
+            result, probs = get_predictions_comment(comment)
+        st.success(result)
+        st.write("Probabilités:")
+        neutral_prob, palestine_prob, israel_prob = probs
+        st.slider("Neutre", 0.0, 1.0, neutral_prob, key="neutral_slider")
+        st.slider("Avec Palestine", 0.0, 1.0, palestine_prob, key="palestine_slider")
+        st.slider("Avec Israël", 0.0, 1.0, israel_prob, key="israel_slider")
+
 if __name__ == "__main__":
     main()