Update app.py

app.py

@@ -1,7 +1,3 @@
-# ==============================
-# INSTALL REQUIRED PACKAGES
-# ==============================
-
 # ==============================
 # IMPORTS
 # ==============================
@@ -9,44 +5,55 @@ import torch
 import re
 import time
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel
+import zipfile
+import os
 
 # ==============================
-# GPU STATUS FUNCTION
+# UNZIP MODEL
+# ==============================
+zip_path = "./1epoch-cds-mistral4bit-training.zip"
+extract_path = "./1epoch-cds-mistral4bit-training"
+
+if not os.path.exists(extract_path):
+    print("📦 Extracting model...")
+    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
+        zip_ref.extractall(extract_path)
+    print("✅ Done")
+else:
+    print("✅ Model already extracted")
+
+# ==============================
+# GPU STATUS FUNCTION (SAFE)
 # ==============================
-import os
 def print_gpu():
-    os.system("nvidia-smi")
+    try:
+        os.system("nvidia-smi")
+    except:
+        print("No GPU (running on CPU)")
 
 # ==============================
-# LOAD MODEL + LO-RA ADAPTER
+# LOAD MODEL
 # ==============================
 BASE_MODEL = "mistralai/Mistral-7B-v0.1"
-MODEL_PATH = "/kaggle/input/datasets/malaikaahmed971/1epoch-cds-mistral4bit-training/mistral7b_fast/checkpoint-7125"
-
-print("🚀 Loading base model and LoRA adapter...")
+MODEL_PATH = "./1epoch-cds-mistral4bit-training/mistral7b_fast/checkpoint-7125"
 
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_compute_dtype=torch.float16,
-    bnb_4bit_use_double_quant=True,
-    bnb_4bit_quant_type="nf4"
-)
+print("🚀 Loading model...")
 
 tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
 
+# CPU SAFE LOADING
 base_model = AutoModelForCausalLM.from_pretrained(
     BASE_MODEL,
-    quantization_config=bnb_config,
-    device_map="auto"
+    device_map={"": "cpu"}
 )
 
 model = PeftModel.from_pretrained(base_model, MODEL_PATH)
 model.eval()
 
 # ==============================
-# CLEAN OUTPUT & METRICS FUNCTIONS
+# CLEAN OUTPUT
 # ==============================
 def clean_output(text):
     text = text.strip()
@@ -57,6 +64,9 @@ def clean_output(text):
         return before.strip() + "\n\nFinal Decision: " + decision_line.strip()
     return text
 
+# ==============================
+# LABEL EXTRACTION
+# ==============================
 def extract_label(text):
     text = text.lower()
     if "final decision" in text:
@@ -67,54 +77,74 @@ def extract_label(text):
             return "hate"
     return "unknown"
 
+# ==============================
+# METRICS
+# ==============================
 def compute_metrics(output, post):
     out = output.lower()
     post = post.lower()
+
     htc = 1 if "final decision" in out else 0
+
     post_words = post.split()
     qf = 1 if any(word in out for word in post_words[:5]) else 0
+
     tgi_keywords = ["muslim","black","white","asian","women","men","jews","christian","pakistan","indian"]
     tgi = 1 if any(word in out for word in tgi_keywords) else 0
+
     pred = extract_label(out)
-    cons = 1 if (pred=="hate" and "hate" in out) or (pred=="non_hate" and "not hate" in out) else 0
+
+    if pred == "hate" and "hate" in out:
+        cons = 1
+    elif pred == "non_hate" and "not hate" in out:
+        cons = 1
+    else:
+        cons = 0
+
     return htc, qf, tgi, cons
 
 # ==============================
-# INFERENCE FUNCTION FOR GRADIO
+# INFERENCE FUNCTION
 # ==============================
 def infer(post):
-    print("\n📊 GPU STATUS BEFORE INFERENCE:")
+    print("\n📊 STATUS BEFORE:")
     print_gpu()
-    
-    inputs = tokenizer(post, return_tensors="pt").to(model.device)
+
+    inputs = tokenizer(post, return_tensors="pt")
     token_count = inputs.input_ids.shape[1]
+
     start = time.time()
-    
+
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=300,
+            max_new_tokens=300,   # SAME as your original
             min_new_tokens=150,
             do_sample=False,
             repetition_penalty=1.15,
             pad_token_id=tokenizer.eos_token_id
         )
-    
+
     end = time.time()
+
     output = tokenizer.decode(outputs[0], skip_special_tokens=True)
     response = clean_output(output.replace(post, "").strip())
+
     pred = extract_label(response)
     htc, qf, tgi, cons = compute_metrics(response, post)
-    
-    metrics = f"Prediction: {pred}\nHTC={htc}, QF={qf}, TGI={tgi}, Consistency={cons}\nTokens processed: {token_count}\nTime: {end-start:.2f} sec"
-    
-    print("\n📊 GPU STATUS AFTER INFERENCE:")
+
+    metrics = f"""Prediction: {pred}
+HTC={htc}, QF={qf}, TGI={tgi}, Consistency={cons}
+Tokens processed: {token_count}
+Time: {end-start:.2f} sec"""
+
+    print("\n📊 STATUS AFTER:")
     print_gpu()
-    
+
     return response, metrics
 
 # ==============================
-# CREATE GRADIO INTERFACE
+# GRADIO UI
 # ==============================
 iface = gr.Interface(
     fn=infer,
@@ -124,8 +154,10 @@ iface = gr.Interface(
         gr.Textbox(label="Metrics", lines=5)
     ],
     title="Hate Speech Rationales + Decision",
-    description="Enter a social media post, and the model will generate step-by-step rationales and the final decision."
+    description="Enter a post to get step-by-step reasoning + final decision"
 )
 
-# Launch the app with public share link
-iface.launch(share=True)
+# ==============================
+# LAUNCH
+# ==============================
+iface.launch()