Update app.py

app.py

@@ -29,6 +29,7 @@ import os
 from flask import Flask, jsonify, request
 import requests
 from fetch_data import fetch_and_update_training_data
+import gradio as gr
 
 # Load configuration file 
 with open('config.json', 'r') as config_file:
@@ -308,3 +309,56 @@ else:
     model = AutoModelForSequenceClassification.from_pretrained(model_save_path).to('cpu')
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_save_path)
 
+#Function to classify user input
+def classify_user_input(user_input):
+    while True:
+           
+        # 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")
+
+
+iface = gr.Interface(fn=classify_user_input, inputs="text", outputs="text") 
+iface.launch(share=True)