Update app.py

app.py

@@ -103,6 +103,9 @@ codes = """001 - Pendaftaran Kendaraan
 008 - Pelaporan dan Bantuan Terkait Kendaraan
 009 - Penangguhan atau Deklarasi Perubahan Penggunaan Kendaraan""".split("\n")
 
+codes = """003 - Alih Kepemilikan atau Balik Nama Kendaraan
+005 - Pembayaran dan Pengelolaan Pajak Kendaraan""".split("\n")
+
 # documents = """Vehicle Registration
 # Vehicle Data Update
 # Vehicle Ownership or Change of Name
@@ -416,8 +419,8 @@ def load_model(model_id):
         return model, None
     else:
         model = SentenceTransformer(model_id, trust_remote_code=True)
-        # codes_emb = model.encode([x[6:] for x in codes])
-        codes_emb = model.encode([x["examples"] for x in examples])#.mean(axis=1)
+        codes_emb = model.encode([x[6:] for x in codes])
+        # codes_emb = model.encode([x["examples"] for x in examples])#.mean(axis=1)
         # codes_emb = np.mean([model.encode(x["examples"]) for x in examples], axis=1)
         return model, codes_emb
         
@@ -740,9 +743,20 @@ def reload(chosen_model_id):
     return f"Model {chosen_model_id} is ready!"
     
 
-def respond_pkb(message, history: list[tuple[str, str]]):
-    tarif_pkb = 1.5/100
-    tarif_bbnkb = 0.1
+def respond_pkb(
+    message, 
+    history: list[tuple[str, str]], 
+    threshold,
+    tarif_pkb,
+    tarif_bbnkb,
+):
+    global codes_emb
+    global undetected
+
+    undetected_code = undetected[:3]
+    
+    # tarif_pkb = 1.5/100
+    # tarif_bbnkb = 0.1
 
     pattern = r'\b([A-Za-z]{1,2})\s?(\d{2,4})\s?([A-Za-z]{1,3})\b'
     
@@ -750,6 +764,19 @@ def respond_pkb(message, history: list[tuple[str, str]]):
     
     plates = [" ".join(x).upper() for i,x in enumerate(matches)]
 
+
+    if type(model) == CrossEncoder:
+        sentence_pairs = [[message, v["name"]] for v in detail_perhitungan.values()]
+        scores = model.predict(sentence_pairs, convert_to_tensor=True)
+        # weights = [1,1]
+    else:
+        text_emb = model.encode(message)
+        scores = cos_sim(codes_emb, text_emb).mean(axis=-1)#[:,0]
+        # weights = [19,8]
+        
+    # scores_argsort = scores.argsort(descending=True)
+
+
     out = ""
     
     for inp in plates:
@@ -772,10 +799,24 @@ def respond_pkb(message, history: list[tuple[str, str]]):
           for k,v in vehicle.iloc[0].items():
             out += f"\n{k} \t\t: {v}"
           # out += "\n=================================================================="
-          out += "\n--"
-          out += f"\nBea Balik Nama Kendaraan Bermotor (BBNKB) : Rp{int(bbnkb):,}"
-          out += f"\nPajak Kendaraan Bermotor (PKB) : Rp{int(pkb):,}\n"
             
+          out += "\n--"
+
+          if scores[0] > threshold:
+              out += "\nRumus Bea Balik Nama Kendaraan Bermotor (BBNKB) : TARIF x NJKB x PENGENAAN"
+              out += f"\nKalkulasi : 10% x Rp{int(nilai_jual):,} x 100%"
+              out += f"\nTotal Pembayaran : Rp{int(bbnkb):,}\n"
+          if scores[1] > threshold:
+              out += "\nRumus Pokok Pajak Kendaraan Bermotor (PKB) : TARIF * NJKB * BOBOT * PENGENAAN"
+              out += f"\nKalkulasi : 1% * Rp{int(nilai_jual):,} * {bobot} * 100%"
+              out += f"\nTotal Pembayaran : Rp{int(pkb):,}\n"
+
+    if not out:
+        if scores[0] > threshold:
+            out += "\nRumus Bea Balik Nama Kendaraan Bermotor (BBNKB) : TARIF x NJKB x PENGENAAN"
+        if scores[1] > threshold:
+            out += "\nRumus Pokok Pajak Kendaraan Bermotor (PKB) : TARIF * NJKB * BOBOT * PENGENAAN"
+
     return out
 
 
@@ -810,11 +851,13 @@ with gr.Blocks() as demo:
     chat_interface = gr.ChatInterface(
         respond_pkb, 
         chatbot=gr.Chatbot(height=670),#800
-        # additional_inputs=[
-        #     gr.Number(0.44, label="confidence threshold", show_label=True, minimum=0., maximum=1.0, step=0.1),
+        additional_inputs=[
+            gr.Number(0.25, label="confidence threshold", show_label=True, minimum=0., maximum=1.0, step=0.1),
+            gr.Number(0.015, label="tarif pajak kendaraan bermotor (pkb)", show_label=True, minimum=0., maximum=1.0, step=0.01),
+            gr.Number(0.1, label="tarif bea balik nama kendaraan bermotor (bbnkb)", show_label=True, minimum=0., maximum=1.0, step=0.05),
         #     gr.Checkbox(True, label="multiple", info="Allow multiple request code numbers"),
         #     gr.Number(3, label="maximum number of request codes", show_label=True, minimum=0, step=1, precision=0)
-        # ],
+        ],
         # # type="messages",
         # # textbox=gr.Textbox(placeholder="Ask me a yes or no question", container=False, scale=7),
         # # title="SamSat Virtual Assistant",