.

README.md

@@ -14,4 +14,13 @@ license: mit
 short_description: 'codebert-base-mlm: a fill-in-middle/masked language model'
 ---
 
-An example chatbot using [Gradio](https://gradio.app), [`huggingface_hub`](https://huggingface.co/docs/huggingface_hub/v0.22.2/en/index), and the [Hugging Face Inference API](https://huggingface.co/docs/api-inference/index).
+A CodeBERT Masked Language Model demo using [Gradio](https://gradio.app) and [Transformers](https://huggingface.co/docs/transformers). This app predicts masked tokens in code snippets.
+
+## Usage
+
+Enter code with `<mask>` tokens where you want predictions:
+- `def <mask>(x, y): return x + y`
+- `import <mask>`
+- `for i in <mask>(10):`
+
+The model will suggest the most likely tokens to fill in the mask.

app.py

@@ -7,43 +7,65 @@ model_name = "microsoft/codebert-base-mlm"
 tokenizer = RobertaTokenizer.from_pretrained(model_name)
 model = RobertaForMaskedLM.from_pretrained(model_name)
 
-def predict_masked_code(code_with_mask, top_k=5):
+def predict(code, num_predictions=5):
     """
     Predict the masked token in code.
     Use <mask> to indicate where to predict.
+    
+    Args:
+        code: Code snippet with <mask> token
+        num_predictions: Number of top predictions to return
+    
+    Returns:
+        JSON object with predictions
     """
     try:
         # Replace <mask> with the tokenizer's mask token
-        code_with_mask = code_with_mask.replace("<mask>", tokenizer.mask_token)
+        code_input = code.replace("<mask>", tokenizer.mask_token)
         
         # Tokenize input
-        inputs = tokenizer(code_with_mask, return_tensors="pt")
+        inputs = tokenizer(code_input, return_tensors="pt")
         
         # Find the position of the mask token
         mask_token_index = torch.where(inputs["input_ids"] == tokenizer.mask_token_id)[1]
         
         if len(mask_token_index) == 0:
-            return "Error: No <mask> token found in the input. Please include <mask> where you want predictions."
+            return {
+                "error": "No <mask> token found in the input. Please include <mask> where you want predictions.",
+                "predictions": []
+            }
         
         # Get predictions
         with torch.no_grad():
             outputs = model(**inputs)
-            predictions = outputs.logits
+            logits = outputs.logits
         
         # Get top-k predictions for the mask token
-        mask_token_logits = predictions[0, mask_token_index, :]
-        top_tokens = torch.topk(mask_token_logits, top_k, dim=1)
+        mask_token_logits = logits[0, mask_token_index, :]
+        top_tokens = torch.topk(mask_token_logits, num_predictions, dim=1)
         
-        results = []
-        for i, (token_id, score) in enumerate(zip(top_tokens.indices[0].tolist(), top_tokens.values[0].tolist())):
+        predictions = []
+        for rank, (token_id, score) in enumerate(zip(top_tokens.indices[0].tolist(), top_tokens.values[0].tolist()), 1):
             predicted_token = tokenizer.decode([token_id])
-            filled_code = code_with_mask.replace(tokenizer.mask_token, predicted_token)
-            results.append(f"{i+1}. {predicted_token} (score: {score:.2f})\n   Code: {filled_code}")
+            completed_code = code_input.replace(tokenizer.mask_token, predicted_token)
+            
+            predictions.append({
+                "rank": rank,
+                "token": predicted_token,
+                "score": round(float(score), 4),
+                "completed_code": completed_code
+            })
         
-        return "\n\n".join(results)
+        return {
+            "original_code": code,
+            "predictions": predictions
+        }
     
     except Exception as e:
-        return f"Error: {str(e)}"
+        return {
+            "error": str(e),
+            "predictions": []
+        }
 
 # Create Gradio interface
 with gr.Blocks(title="CodeBERT Masked Language Model") as demo:
@@ -69,7 +91,7 @@ with gr.Blocks(title="CodeBERT Masked Language Model") as demo:
                 lines=5,
                 value="def <mask>(x, y):\n    return x + y"
             )
-            top_k_slider = gr.Slider(
+            num_predictions_slider = gr.Slider(
                 minimum=1,
                 maximum=10,
                 value=5,
@@ -79,10 +101,8 @@ with gr.Blocks(title="CodeBERT Masked Language Model") as demo:
             predict_btn = gr.Button("Predict", variant="primary")
         
         with gr.Column():
-            output = gr.Textbox(
-                label="Predictions",
-                lines=15,
-                interactive=False
+            output = gr.JSON(
+                label="Predictions"
             )
     
     # Examples
@@ -95,12 +115,12 @@ with gr.Blocks(title="CodeBERT Masked Language Model") as demo:
             ["if x <mask> 0:", 5],
             ["class <mask>:", 5],
         ],
-        inputs=[code_input, top_k_slider],
+        inputs=[code_input, num_predictions_slider],
     )
     
     predict_btn.click(
-        fn=predict_masked_code,
-        inputs=[code_input, top_k_slider],
+        fn=predict,
+        inputs=[code_input, num_predictions_slider],
         outputs=output
     )