Update app.py

app.py

@@ -2,56 +2,55 @@ import os
 import gradio as gr
 from transformers import pipeline
 
-model_id = "Badhon/Bangla_punctuation_restore"
+MODEL_ID = "Badhon/Bangla_punctuation_restore"
+
+LABEL_TO_PUNCT = {
+    "COMMA": "،",
+    "DARI": "।",
+    "QUESTION": "?",
+    "EXCLAMATION": "!",
+    "SEMICOLON": ";",
+    "COLON": ":",
+    "HYPHEN": "-"
+}
 
-# Token classification pipeline
 punctuator = pipeline(
     "token-classification",
-    model=model_id,
+    model=MODEL_ID,
     aggregation_strategy="simple",
-    token=os.getenv("HF_TOKEN")  # required if model is private
+    token=os.getenv("HF_TOKEN")
 )
 
-def restore_punctuation(text):
+def restore_punctuation(text: str) -> str:
     if not text.strip():
         return ""
 
-    tokens = punctuator(text)
-    words = text.split()
+    preds = punctuator(text)
 
-    output = []
-    word_idx = 0
+    output = text
+    offset = 0
 
-    for token in tokens:
-        word = words[word_idx]
-        label = token["entity_group"]
+    for p in preds:
+        label = p["entity_group"]
+        if label == "O":
+            continue
 
-        output.append(word)
+        punct = LABEL_TO_PUNCT.get(label)
+        if not punct:
+            continue
 
-        if label == "COMMA":
-            output.append("，")
-        elif label == "PERIOD":
-            output.append("।")
-        elif label == "QUESTION":
-            output.append("?")
-        elif label == "EXCLAMATION":
-            output.append("!")
-        elif label == "COLON":
-            output.append(":")
-        elif label == "SEMICOLON":
-            output.append(";")
+        end = p["end"] + offset
+        output = output[:end] + punct + output[end:]
+        offset += len(punct)
 
-        word_idx += 1
+    return output
 
-    return "".join(output)
-
-# Gradio UI
 demo = gr.Interface(
     fn=restore_punctuation,
-    inputs=gr.Textbox(lines=4, placeholder="Enter Bangla text without punctuation"),
+    inputs=gr.Textbox(lines=4, placeholder="বাংলা টেক্সট লিখুন (যতিচিহ্ন ছাড়া)"),
     outputs="text",
-    title="Bangla Punctuation Restoration (BERT)",
-    description="Bangla punctuation restoration using sagor-bert-base token classification"
+    title="Bangla Punctuation Restoration",
+    description="sagor-bert-base based Bangla punctuation restoration"
 )
 
 demo.launch()