Update app.py

app.py

@@ -6,17 +6,13 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 MODEL_ID = "PraneetNS/codesentinel-full"
 TOKENIZER_ID = "PraneetNS/codesentinel-adapter"
 
+print("Loading tokenizer...")
+
 tokenizer = AutoTokenizer.from_pretrained(
     TOKENIZER_ID,
     trust_remote_code=True,
 )
 
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    dtype=torch.float16,
-    device_map="auto",
-    trust_remote_code=True,
-)
 SYSTEM_PROMPT = """You are CodeSentinel.
 
 You are an expert AI software engineer.
@@ -33,12 +29,10 @@ Capabilities:
 Never generate malware or unsafe code.
 """
 
-# Global model variable
 model = None
 
 
 def load_model():
-    """Load the model only once after GPU allocation."""
     global model
 
     if model is None:
@@ -67,9 +61,30 @@ def generate(message, history):
         }
     ]
 
+    # Convert Gradio history
     if history:
-        for msg in history:
-            messages.append(msg)
+        for item in history:
+            if isinstance(item, (list, tuple)) and len(item) == 2:
+                user_msg, assistant_msg = item
+
+                if user_msg:
+                    messages.append(
+                        {
+                            "role": "user",
+                            "content": user_msg,
+                        }
+                    )
+
+                if assistant_msg:
+                    messages.append(
+                        {
+                            "role": "assistant",
+                            "content": assistant_msg,
+                        }
+                    )
+
+            elif isinstance(item, dict):
+                messages.append(item)
 
     messages.append(
         {
@@ -82,13 +97,18 @@ def generate(message, history):
         messages,
         tokenize=True,
         add_generation_prompt=True,
-        return_tensors="pt",
         return_dict=True,
-    ).to(model.device)
+        return_tensors="pt",
+    )
+
+    inputs = {
+        k: v.to(model.device)
+        for k, v in inputs.items()
+    }
 
     with torch.no_grad():
         outputs = model.generate(
-            inputs,
+            **inputs,
             max_new_tokens=512,
             temperature=0.2,
             top_p=0.95,
@@ -97,8 +117,10 @@ def generate(message, history):
             pad_token_id=tokenizer.eos_token_id,
         )
 
+    generated_tokens = outputs[0][inputs["input_ids"].shape[-1]:]
+
     response = tokenizer.decode(
-        outputs[0][inputs.shape[-1]:],
+        generated_tokens,
         skip_special_tokens=True,
     )
 
@@ -114,7 +136,7 @@ demo = gr.ChatInterface(
         "Explain this C++ function.",
         "Optimize this SQL query.",
         "Write a secure FastAPI login API.",
-        "Review this Java code for security vulnerabilities."
+        "Review this Java code for security vulnerabilities.",
     ],
 )