Update app.py

app.py

@@ -297,7 +297,8 @@ if (should_train_model=='1'): #train model
         token=api_token,
         commit_message="Push model and tokenizer",
     )
-        
+
+    print("Operation complete for fine-tunning.")
 else:
     print('Load Pre-trained')
     model_save_path = f"./{model_save_path}"
@@ -307,65 +308,3 @@ else:
     model = AutoModelForSequenceClassification.from_pretrained(model_save_path).to('cpu')
     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()