Update app.py

app.py

@@ -20,31 +20,44 @@ config = AutoConfig.from_pretrained(model_name, revision=revision, trust_remote_
 print(f"Loading model from {model_name}...")
 model = CustomModel.from_pretrained(model_name, config=config, revision=revision, trust_remote_code=True)
 
+# Check if the model loaded successfully
 if model is None:
     print("Failed to load model. Exiting...")
     exit(1)
 else:
     print("Model loaded successfully.")
 
+# Define the text classification function
 def classify_text(text):
-    inputs = tokenizer(text, return_tensors="pt")
-    outputs = model(**inputs)
-    logits = outputs.logits
-    probabilities = torch.softmax(logits, dim=-1).tolist()[0]
-    predicted_class = torch.argmax(logits, dim=-1).item()
-    return {
-        "Predicted Class": predicted_class,
-        "Probabilities": probabilities
-    }
+    try:
+        # Tokenize the input text
+        inputs = tokenizer(text, return_tensors="pt")
+        # Pass the inputs to the model
+        outputs = model(**inputs)
+        # Get the logits and probabilities
+        logits = outputs.logits
+        probabilities = torch.softmax(logits, dim=-1).tolist()[0]
+        # Get the predicted class
+        predicted_class = torch.argmax(logits, dim=-1).item()
+        return {
+            "Predicted Class": predicted_class,
+            "Probabilities": probabilities
+        }
+    except Exception as e:
+        print(f"Error during text classification: {e}")
+        return {
+            "Predicted Class": "Error",
+            "Probabilities": []
+        }
 
 # Create a Gradio interface
 try:
     iface = gr.Interface(
-        fn=classify_text,
-        inputs=gr.Textbox(lines=2, placeholder="Enter text here..."),
+        fn=classify_text,  # Function to call
+        inputs=gr.Textbox(lines=2, placeholder="Enter text here..."),  # Input component
         outputs=[
-            gr.Label(label="Predicted Class"),
-            gr.Label(label="Probabilities")
+            gr.Label(label="Predicted Class"),  # Output component for predicted class
+            gr.Label(label="Probabilities")  # Output component for probabilities
         ],
         title="DeepSeek-V3 Text Classification",
         description="Classify text using the DeepSeek-V3 model."