update app.py

app.py

@@ -174,7 +174,8 @@ model = LSGMBartForConditionalGeneration.from_pretrained(path, use_auth_token=au
 tokenizer = AutoTokenizer.from_pretrained(path, use_auth_token=auth_token)
 pipe = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
 
-def sim(question, num_beams, text):
+
+def sim(question, X, num_beams, text):
     prefix = "[Frais irrépétibles ; article 700]"
     q_type = int(question[0])
     num_beams = None if num_beams == 0 else num_beams
@@ -182,6 +183,11 @@ def sim(question, num_beams, text):
         question = "Qui sont les demandeurs initiaux (ceux qui assignent) ?"
     elif q_type == 1:
         question = "Qui sont les demandeurs, le résultat de chaque demande et les quanta obtenus?"
+    elif q_type == 2:
+        if X == "":
+            return "Aucune partie renseignée."
+        prefix = "[Frais irrépétibles ; article 700" + X + "]"
+        question = "Combien a obtenu ce demandeur et quel est le résultat de sa demande ?"
     
     text = prefix + " " + question + " <sep> " + text
     generated_text = pipe(
@@ -197,25 +203,38 @@ def sim(question, num_beams, text):
     generated_tokens = generated_text[0]["generated_token_ids"]
     tokens = generated_tokens.tolist()
     sentence = tokenizer.decode(tokens).replace("<s>", "").replace("</s>", "").strip()
-    
-    if q_type == 0:
-        sentence = sentence.split("<liste>")
-        return "Demandeurs initiaux :\n" + "\n".join(sentence)
-    elif q_type == 1:
+
+
+    def process_inputs(sentence, labels):
+        labels = [i + " : \n" for i in labels]
         demandes = sentence.split("<demande>")
         final_results = []
         for i, demande in enumerate(demandes):
             results = []
             elems = demande.split("<sep>")
-            for prefix, elem in zip(["Demandeurs : \n", "Résultat : \n", "Montant obtenu : \n"], elems):
+            for prefix, elem in zip(labels, elems):
                 results.append(prefix + "\n".join(elem.split("<liste>")))
             final_results.append("Demande " + str(i+1) + "\n\n" +"\n".join(results))
         return "\n\n".join(final_results)
 
+    if q_type == 0:
+        sentence = sentence.split("<liste>")
+        return "Demandeurs initiaux :\n" + "\n".join(sentence)
+    elif q_type == 1:
+        return process_inputs(sentence, ["Demandeurs", "Résultat", "Montant obtenu"])
+    elif q_type == 2:
+        return process_inputs(sentence, ["Montant obtenu", "Résultat"])
+
 iface = gr.Interface(
     sim, 
     inputs=[
-        gr.Dropdown(["0: Demandeur initial", "1: Récapitulatif des demandes"], label="Information à générer", info="Succeptible de changer."), 
+        gr.Dropdown([
+            "0: Demandeur initial", 
+            "1: Récapitulatif des demandes", 
+            "2: Récapitulatif des demandes pour une partie"], 
+            value="1: Récapitulatif des demandes",
+            label="Information à générer", info="Succeptible de changer."),
+        gr.Text(placeholder="Exemple: X..."), 
         gr.Slider(label="Nombre de Beams", minimum=0, maximum=10, value=5, step=1), 
         gr.Textbox(label= "Décision de justice", placeholder="Décision à traiter...", lines=25)
         ], 
@@ -225,7 +244,7 @@ iface = gr.Interface(
     allow_flagging="auto", 
     title="Extraction",
     description="Extrait les demandes relatives à l'article 700",
-    examples=[["Récapitulatif des demandes", 5, example_1]],
+    examples=[["Récapitulatif des demandes", "", 5, example_1]],
     #live=True,
     )
 iface.launch()