Update app.py

app.py

@@ -1,24 +1,48 @@
 import gradio as gr
-from transformers import pipeline
+import torch
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
-# This is the key line:
-# It loads your private model from your *other* repository.
-pipe = pipeline(
-    "text-classification",
-    model="kxshrx/infrnce-bert-classifier"
+# --- Configuration ---
+MODEL_NAME = "distilbert-base-uncased"
+NUM_LABELS = 6
+MODEL_PATH = "controlled_bert_model.pth" # The name of the file you uploaded
+
+# --- Load Tokenizer and Model ---
+print("Loading tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+
+print("Loading model architecture...")
+# First, create the model "shell"
+model = AutoModelForSequenceClassification.from_pretrained(
+    MODEL_NAME, 
+    num_labels=NUM_LABELS
+)
+
+print(f"Loading fine-tuned weights from {MODEL_PATH}...")
+# Now, load your trained weights into the shell
+model.load_state_dict(
+    torch.load(MODEL_PATH, map_location=torch.device("cpu"))
 )
+model.eval() # Set model to evaluation mode
+print("Model loaded successfully!")
 
 def classify_log(log_text):
-    # This function runs the classification.
-    results = pipe(log_text, top_k=None)
-    # We format the result into a simple dictionary.
-    return {item['label']: item['score'] for item in results[0]}
+    """
+    This function runs the classification using your loaded .pth model.
+    """
+    inputs = tokenizer(log_text, return_tensors="pt", padding=True, truncation=True)
+    with torch.no_grad():
+        logits = model(**inputs).logits
+    
+    scores = torch.softmax(logits, dim=1).squeeze().tolist()
+    # Create a dictionary of {label_name: score}
+    confidences = {model.config.id2label[i]: score for i, score in enumerate(scores)}
+    return confidences
 
-# This creates a simple web UI for testing and, more importantly,
-# an API endpoint that we can call.
+# This creates the Gradio interface and API endpoint
 gr.Interface(
     fn=classify_log,
     inputs=gr.Textbox(lines=5, label="Log Entry"),
-    outputs=gr.Label(num_top_classes=6),
+    outputs=gr.Label(num_top_classes=6, label="Classification Results"),
     title="Infrnce Private Log Classifier API"
 ).launch()