Update app.py

app.py

@@ -2,11 +2,11 @@ import torch
 from transformers import BertTokenizerFast, BertForTokenClassification
 import gradio as gr
 
-# Load tokenizer and model
+# Initialize tokenizer and model
 tokenizer = BertTokenizerFast.from_pretrained('bert-base-uncased')
 model = BertForTokenClassification.from_pretrained('maximuspowers/bias-detection-ner')
-model.eval()  # Set the model to evaluation mode
-model.to('cuda' if torch.cuda.is_available() else 'cpu')  # Move model to appropriate device
+model.eval()
+model.to('cuda' if torch.cuda.is_available() else 'cpu')
 
 # Define label mappings
 id2label = {
@@ -20,12 +20,10 @@ id2label = {
 }
 
 def predict_ner_tags(sentence):
-    # Tokenize the sentence
     inputs = tokenizer(sentence, return_tensors="pt", padding=True, truncation=True, max_length=128)
     input_ids = inputs['input_ids'].to(model.device)
     attention_mask = inputs['attention_mask'].to(model.device)
 
-    # Predict using the model
     with torch.no_grad():
         outputs = model(input_ids=input_ids, attention_mask=attention_mask)
         logits = outputs.logits
@@ -43,7 +41,6 @@ def predict_ner_tags(sentence):
     return result
 
 def format_output(result):
-    # Create HTML content with formatted output
     formatted_output = "<div style='font-family: Arial;'>"
     for token, labels in result:
         styles = []
@@ -62,12 +59,11 @@ def format_output(result):
 iface = gr.Interface(
     fn=predict_ner_tags,
     inputs="text",
-    outputs=gr.outputs.HTML(label="Output"),
+    outputs="html",  # Directly use "html" here
     title="Named Entity Recognition with BERT",
     description="Enter a sentence to predict NER tags using a BERT model trained for multi-label classification. Different styles represent different entity types.",
     examples=["Tall men are so clumsy."],
-    allow_flagging="never",
-    theme="default"
+    allow_flagging="never"
 )
 
 if __name__ == "__main__":