Update app.py

app.py

@@ -1,15 +1,12 @@
 import gradio as gr
 from transformers import pipeline
-import torch
 
-# Use CPU on Hugging Face Spaces free tier
 DEVICE = -1
-
-# Lazy loading (loads model only when first used)
 models = {}
 
 def get_model(task_name):
     if task_name not in models:
+
         if task_name == "Chatbot":
             models[task_name] = pipeline(
                 "text-generation",
@@ -46,6 +43,13 @@ def get_model(task_name):
                 device=DEVICE
             )
 
+        elif task_name == "Fill Mask":
+            models[task_name] = pipeline(
+                "fill-mask",
+                model="bert-base-uncased",
+                device=DEVICE
+            )
+
     return models[task_name]
 
 
@@ -62,8 +66,7 @@ def run_task(task, user_input, chat_history):
             max_new_tokens=100,
             pad_token_id=50256
         )
-
-        bot_reply = response[0]["generated_text"]
+        bot_reply = response[0]["generated_text"][len(user_input):]
         chat_history = chat_history + [(user_input, bot_reply)]
         return "", chat_history
 
@@ -85,7 +88,6 @@ def run_task(task, user_input, chat_history):
         entities = model(user_input)
         if not entities:
             return "No entities found.", chat_history
-
         formatted = "\n".join(
             f"{e['word']} ({e['entity_group']}) - {e['score']:.2f}"
             for e in entities
@@ -96,6 +98,16 @@ def run_task(task, user_input, chat_history):
         translation = model(user_input)[0]["translation_text"]
         return translation, chat_history
 
+    elif task == "Fill Mask":
+        if "<mask>" not in user_input:
+            return "Please include the token <mask> in your sentence.", chat_history
+        predictions = model(user_input)
+        formatted = "\n".join(
+            f"{p['token_str']} (score: {p['score']:.4f})"
+            for p in predictions
+        )
+        return formatted, chat_history
+
 
 with gr.Blocks(title="NLP Application") as demo:
 
@@ -107,7 +119,8 @@ with gr.Blocks(title="NLP Application") as demo:
             "Sentiment Analysis",
             "NER",
             "Summarization",
-            "Translation (EN→FR)"
+            "Translation (EN→FR)",
+            "Fill Mask"
         ],
         label="Select NLP Task"
     )
@@ -118,9 +131,7 @@ with gr.Blocks(title="NLP Application") as demo:
         label="Input Text"
     )
 
-    output_box = gr.Textbox(
-        label="Output"
-    )
+    output_box = gr.Textbox(label="Output")
 
     chatbot = gr.Chatbot(label="Conversation")
 
@@ -134,4 +145,6 @@ with gr.Blocks(title="NLP Application") as demo:
         outputs=[output_box, chatbot]
     )
 
-demo.launch()
+
+if __name__ == "__main__":
+    demo.launch()