Update app.py

app.py

@@ -308,3 +308,64 @@ else:
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_save_path)
 
 
+#Define the label mappings (this must match the mapping used during training)
+label_mapping = model.config.label_mapping
+label_mapping_reverse = {value: key for key, value in label_mapping.items()}
+
+
+#Function to classify user input
+def classify_user_input():
+    while True:
+        user_input = input("Enter a command (or type 'q' to quit): ")
+        if user_input.lower() == 'q':
+            print("Exiting...")
+            break
+             
+        # Tokenize and predict
+        input_encoding = tokenizer(user_input, padding=True, truncation=True, return_tensors="pt").to('cuda')
+        
+        with torch.no_grad():
+            #attention_mask = input_encoding['attention_mask'].clone()
+
+            # Modify the attention mask to emphasize certain key tokens
+            for idx, token_id in enumerate(input_encoding['input_ids'][0]):
+                word = tokenizer.decode([token_id])
+                print(word)
+                #if word.strip() in ["point", "summarize", "oil", "maintenance"]:  # Target key tokens
+                    #attention_mask[0, idx] = 2  # Increase attention weight for these words
+                # else:
+                #     attention_mask[0, idx] = 0 
+            #print (attention_mask)
+            #input_encoding['attention_mask'] = attention_mask               
+            output = model(**input_encoding, output_hidden_states=True)
+            # print('start-logits')
+            # print(output.logits)
+            # print('end-logits')
+            #print(output)
+            attention = output.attentions  # Get attention scores
+            #print('atten')
+            #print(attention)
+            # Apply softmax to get the probabilities (confidence scores)
+            probabilities = F.softmax(output.logits, dim=-1)
+
+            # tokens = tokenizer.convert_ids_to_tokens(input_encoding['input_ids'][0].cpu().numpy())
+            # # Display the attention visualization
+            # input_text = tokenizer.convert_ids_to_tokens(input_encoding['input_ids'][0])
+
+            prediction = torch.argmax(output.logits, dim=1).cpu().numpy()
+
+            # Map prediction back to label
+            print(prediction)
+            predicted_label = label_mapping_reverse[prediction[0]]
+
+
+            print(f"Predicted intent: {predicted_label}\n")
+            # Print the confidence for each label
+            print("\nLabel Confidence Scores:")
+            for i, label in label_mapping_reverse.items():
+                confidence = probabilities[0][i].item()  # Get confidence score for each label
+                print(f"{label}: {confidence:.4f}")
+            print("\n")
+
+#Run the function
+classify_user_input()