comment-guard-api / verify_model.py
tejesh916K's picture
Deploy: Comment Guard API - FastAPI + MuRIL BERT
b8300d6
"""
MODEL VERIFICATION SCRIPT
Use this to test your trained model locally on your PC.
"""
import os
from transformers import pipeline
def test_model():
# 1. Path to your model folder
# Change this to 'model_output_v2' if testing the new version
model_path = "./model_output"
if not os.path.exists(model_path):
print(f"❌ Error: Model folder '{model_path}' not found.")
print("Please ensure you have moved your Kaggle/Colab output into the 'backend' folder.")
return
print("πŸ”„ Loading model (this may take a few seconds)...")
try:
# Load the toxicity classifier
classifier = pipeline(
"text-classification",
model=model_path,
tokenizer=model_path,
device=-1 # Use -1 for CPU, 0 for first GPU
)
print("βœ… Model loaded successfully!\n")
except Exception as e:
print(f"❌ Failed to load model: {e}")
return
print("Enter 'quit' to exit.")
while True:
text = input("\nπŸ“ Enter a comment to test: ")
if text.lower() == 'quit':
break
if not text.strip():
continue
# Get prediction
result = classifier(text)[0]
label = result['label']
score = result['score']
# Map labels to human-readable text
# LABEL_1 is usually Toxic, LABEL_0 is Safe
is_toxic = "TOXIC πŸ”΄" if label == "LABEL_1" else "SAFE 🟒"
print("-" * 30)
print(f"Result: {is_toxic}")
print(f"Confidence: {score*100:.2f}%")
print("-" * 30)
if __name__ == "__main__":
test_model()