Update app.py

app.py

@@ -13,24 +13,22 @@ id2label = model.config.id2label
 # Inference function
 def classify_text(text):
     if not text.endswith("."):
-        text = text + "."
+        text += "."
     inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
     with torch.no_grad():
         outputs = model(**inputs)
-        probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
-        pred_class = torch.argmax(probs, dim=1).item()
-        confidence = probs[0, pred_class].item()
-    label = id2label[pred_class]
-    return f"Predicted Level: {label} (Confidence: {confidence:.2f})"
-   
+        probs = torch.nn.functional.softmax(outputs.logits, dim=-1)[0]
+    # Create a dictionary of labels and their corresponding probabilities
+    confidences = {id2label[i]: float(probs[i]) for i in range(len(probs))}
+    return confidences
 
 # Gradio interface
 demo = gr.Interface(
     fn=classify_text,
     inputs=gr.Textbox(lines=4, placeholder="Schreibe etwas auf Deutsch..."),
-    outputs=gr.Textbox(label="Language Level Prediction"),
+    outputs=gr.Label(num_top_classes=4),
     title="German Language Level Classifier",
-    description="Enter German text and get the predicted CEFR level (A1 to C2).🚧 Working in progress.🚧"
+    description="Enter German text and get the predicted CEFR level (A1 to C2). 🚧 Work in progress. 🚧"
 )
 
 # Launch app