Update app.py

app.py

@@ -7,6 +7,10 @@ model_name = "TextLabRUET/xlm-r_based_bangla_sentence_classifier"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSequenceClassification.from_pretrained(model_name)
 
+# Set device (GPU if available, otherwise CPU)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+
 # Mapping predicted class to Bangla sentence types
 class_mapping = {
     0: "Assertive Sentence (বর্ণনামূলক বাক্য)",
@@ -17,29 +21,40 @@ class_mapping = {
 }
 
 # Function for prediction
-def predict(sentence):
+def predict_bangla_sentence(sentence):
+    # Tokenize the input sentence
     inputs = tokenizer(sentence, return_tensors="pt", truncation=True, padding=True, max_length=128)
+    
+    # Move input tensors to the same device as the model
+    inputs = {key: val.to(device) for key, val in inputs.items()}
+    
+    # Perform inference
     with torch.no_grad():
         outputs = model(**inputs)
+    
+    # Get the predicted class
     logits = outputs.logits
-    predicted_class = torch.argmax(logits, dim=1).item()
+    predicted_class = torch.argmax(logits, dim=-1).item()
     
+    # Return the predicted sentence type
     sentence_type = class_mapping.get(predicted_class, "Unknown Sentence Type")
     return f"Predicted Class: {sentence_type}"
 
 # Create Gradio UI
 iface = gr.Interface(
-    fn=predict,
+    fn=predict_bangla_sentence,
     inputs=gr.Textbox(lines=2, placeholder="Enter a Bangla sentence..."),
     outputs="text",
     title="Bangla Sentence Classifier",
     description=(
-        "This model was trained on a curated Bangla dataset by **TextLab RUET**. "
-        "It classifies Bangla sentences into five distinct categories: Assertive, Interrogative, Imperative, Optative, and Exclamatory "
-        "using the **XLM-R model**. Enter a Bangla sentence below to see how our model interprets it!"
+        "This model was trained by **TextLabRUET** using an extensive Bangla dataset. "
+        "It classifies Bangla sentences into five types: Assertive, Interrogative, Imperative, Optative, and Exclamatory "
+        "using the **XLM-R model**. Enter a Bangla sentence below to see how our model analyzes it!\n\n"
+        "Note: While we aim for accuracy, the model may occasionally misclassify sentences due to dataset limitations. "
+        "We apologize for any errors and appreciate your understanding." 
     ),
     theme="compact"
 )
 
 # Launch the Gradio app
-iface.launch(share=True)
+iface.launch(share=True)