Update app.py

app.py

@@ -1,42 +1,47 @@
 import gradio as gr
-from transformers import pipeline
+import torch
+from transformers import BertForMaskedLM, BertTokenizer
 
-pipe = pipeline("fill-mask", model="google-bert/bert-base-uncased")
+# Modell und Tokenizer laden
+model_name = "bert-base-uncased"
+model = BertForMaskedLM.from_pretrained(model_name)
+tokenizer = BertTokenizer.from_pretrained(model_name)
 
-title = "BERT"
-description = "Gradio Demo for BERT. To use it, simply add your text, or click one of the examples to load them. Read more at the links below."
-article = "<p style='text-align: center'><a href='https://arxiv.org/abs/1810.04805' target='_blank'>BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding</a></p>"
+# Inferenz-Funktion definieren
+def inference(input_text):
+    if "[MASK]" not in input_text:
+        return "Error: The input text must contain the [MASK] token."
+    
+    # Tokenisierung
+    inputs = tokenizer(input_text, return_tensors="pt")
+    mask_token_index = torch.where(inputs["input_ids"] == tokenizer.mask_token_id)[1]
 
-examples = [
-    ['Paris is the [MASK] of France.', 'bert-base-cased']
-]
+    # Vorhersage
+    with torch.no_grad():
+        outputs = model(**inputs)
+        logits = outputs.logits
 
-# Lade die Interfaces für die Modelle
-io1 = gr.Interface.load("huggingface/bert-base-cased")
-io2 = gr.Interface.load("huggingface/bert-base-uncased")
+    # Wahrscheinlichsten Token für [MASK] finden
+    mask_token_logits = logits[0, mask_token_index, :]
+    top_token = torch.topk(mask_token_logits, 1, dim=1).indices[0].tolist()
 
-def inference(inputtext, model):
-    if "[MASK]" not in inputtext:
-        return {"error": "The input text must contain the [MASK] token."}
+    # Vorhersage in den Text einfügen
+    predicted_token = tokenizer.decode(top_token)
+    result_text = input_text.replace("[MASK]", predicted_token, 1)
     
-    if model == "bert-base-cased":
-        return io1(inputtext)
-    elif model == "bert-base-uncased":
-        return io2(inputtext)
-    else:
-        return {"error": "Invalid model selected"}
+    return result_text
 
+# Gradio Interface definieren
 iface = gr.Interface(
     fn=inference,
-    inputs=[
-        gr.Textbox(label="Context", lines=10, placeholder="Enter text with [MASK] token"),
-        gr.Dropdown(choices=["bert-base-cased", "bert-base-uncased"], value="bert-base-cased", label="model")
-    ],
-    outputs=gr.JSON(label="Output"),  # We use JSON to display errors or outputs
-    examples=examples,
-    article=article,
-    title=title,
-    description=description
+    inputs="text",
+    outputs="text",
+    examples=[
+        ["The capital of France is [MASK]."],
+        ["The quick brown fox jumps over the [MASK] dog."]
+    ]
 )
 
-iface.launch(share=True)
+# Interface starten
+if __name__ == "__main__":
+    iface.launch()