Update app.py

app.py

@@ -22,7 +22,9 @@ def predict_sentiment(text_input, model_selection):
         predicted_class = int(logits.argmax())
         inference_time = end_time - start_time
         model_size = model.num_parameters()
-        return candidate_labels[predicted_class], inference_time, model_size
+        architecture = model.config.architectures[0]
+        #batch_size = inputs['input_ids'].shape[0]
+        return candidate_labels[predicted_class], inference_time, model_size, architecture
     else:
         start_time = time.time()
         result = pretrained_model(text_input, candidate_labels)
@@ -30,7 +32,9 @@ def predict_sentiment(text_input, model_selection):
         predicted_class = result["labels"][0]
         inference_time = end_time - start_time
         model_size = pretrained_tokenizer.model_max_length + pretrained_model.model.num_parameters()
-        return predicted_class, inference_time, model_size
+        architecture = pretrained_model.model.config.architectures[0]
+        #batch_size = 1
+        return predicted_class, inference_time, model_size, architecture
 
 inputs = [
     gr.inputs.Textbox("Enter text"),
@@ -41,6 +45,7 @@ outputs = [
     gr.outputs.Textbox(label="Predicted Sentiment"),
     gr.outputs.Textbox(label="Inference Time (s)"),
     gr.outputs.Textbox(label="Model Size (params)"),
+    gr.outputs.Textbox(label="Architecture"),
 ]
 
 gr.Interface(fn=predict_sentiment, inputs=inputs, outputs=outputs, title="Sentiment Analysis", description="roberta-large-mnli fine tuned with poem_sentiment dataset for sentiment analysis", examples=[