Quickstart updated

README.md

@@ -173,10 +173,12 @@ class TurnDetector:
         self.max_length = 8192
         print("✅ Model and tokenizer loaded successfully.")
 
-    def predict(self, text: str) -> str:
+    def predict(self, text: str) -> tuple:
         """
         Predicts if a given text utterance is the end of a turn.
-        Returns "End of Turn" or "Not End of Turn".
+        Returns (predicted_label, confidence) where:
+        - predicted_label: 0 for "Not End of Turn", 1 for "End of Turn"
+        - confidence: confidence score between 0 and 1
         """
         # Tokenize the input text
         inputs = self.tokenizer(
@@ -194,12 +196,26 @@ class TurnDetector:
         
         # Run inference
         outputs = self.session.run(None, feed_dict)
-        logits = outputs
-        
-        # Get the predicted class (0 or 1)
-        prediction_index = np.argmax(logits, axis=1)
-        
-        return "End of Turn" if prediction_index == 1 else "Not End of Turn"
+        logits = outputs[0]
+
+        probabilities = self._softmax(logits[0])
+        predicted_label = np.argmax(probabilities)
+        print(f"Logits: {logits[0]}")
+        print(f"Probabilities:")
+        print(f"  Not End of Turn: {probabilities[0]:.1%}")
+        print(f"  End of Turn:     {probabilities[1]:.1%}")
+        print(f"Predicted Label: {predicted_label}")
+        print(f"Confidence: {probabilities[predicted_label]:.1%}")
+        print(f"Confidence (class=1): {np.max(probabilities):.1%}")
+        confidence = float(np.max(probabilities))
+
+        return predicted_label, confidence
+
+    def _softmax(self, x, axis=None):
+        if axis is None:
+            axis = -1
+        exp_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
+        return exp_x / np.sum(exp_x, axis=axis, keepdims=True)
 
 # --- Example Usage ---
 if __name__ == "__main__":
@@ -213,9 +229,10 @@ if __name__ == "__main__":
     ]
     
     for sentence in sentences:
-        result = detector.predict(sentence)
-        print(f"'{sentence}' -> {result}")
-
+        predicted_label, confidence = detector.predict(sentence)
+        result = "End of Turn" if predicted_label == 1 else "Not End of Turn"
+        print(f"'{sentence}' -> {result} (confidence: {confidence:.3f})")
+        print("-" * 50)
 ```