Update app.py

app.py

@@ -1,28 +1,45 @@
 import gradio as gr
-from transformers import BertTokenizer, BertModel
+from transformers import pipeline
 
-# Load the model and tokenizer
-tokenizer = BertTokenizer.from_pretrained("google-bert/bert-base-uncased")
-model = BertModel.from_pretrained("google-bert/bert-base-uncased")
+def mask(text):
+  mask_model = pipeline("fill-mask", model="google-bert/bert-base-uncased")
+  output = mask_model(text)
+  return output[0]['sequence']
 
-def predict(input_text):
-    # Preprocess the input text
-    inputs = tokenizer(input_text, return_tensors="pt")
-    # Get the model's output
-    outputs = model(**inputs)
-    # Extract the last hidden state
-    last_hidden_state = outputs.last_hidden_state[:, 0, :]
-    # Return a response based on the output
-    return {"response": last_hidden_state.detach().numpy().tolist()}
+with gr.Blocks() as demo:
+  gr.Markdown("BERT Test - Zircon.tech")
+        with gr.Column(scale=1):
+            textbox_mask = gr.Textbox(label="Give me a sentance and use the word [MASK] to predict what will go there")
+            mask_button = gr.Button("Predict")
+            textbox_mask_output = gr.Textbox()
 
-# Create a Gradio interface
-iface = gr.Interface(
-    fn=predict,
-    inputs="text",
-    outputs="text",
-    title="BERT-based Text Analyzer",
-    description="Enter some text to analyze with BERT!"
-)
+            mask_button.click(fn=mask,
+                          inputs = textbox_mask,
+                          outputs = textbox_mask_output)
+demo.launch(debug=True)
+    
+# # Load the model and tokenizer
+# tokenizer = BertTokenizer.from_pretrained("google-bert/bert-base-uncased")
+# model = BertModel.from_pretrained("google-bert/bert-base-uncased")
 
-# Launch the interface
-iface.launch()
+# def predict(input_text):
+#     # Preprocess the input text
+#     inputs = tokenizer(input_text, return_tensors="pt")
+#     # Get the model's output
+#     outputs = model(**inputs)
+#     # Extract the last hidden state
+#     last_hidden_state = outputs.last_hidden_state[:, 0, :]
+#     # Return a response based on the output
+#     return {"response": last_hidden_state.detach().numpy().tolist()}
+
+# # Create a Gradio interface
+# iface = gr.Interface(
+#     fn=predict,
+#     inputs="text",
+#     outputs="text",
+#     title="BERT-based Text Analyzer",
+#     description="Enter some text to analyze with BERT!"
+# )
+
+# # Launch the interface
+# iface.launch()