Update app.py

app.py

@@ -1,143 +1,66 @@
-import torch
 import gradio as gr
-import spaces
+from huggingface_hub import hf_hub_download
+from llama_cpp import Llama
 
-from transformers import (
-    AutoTokenizer,
-    AutoModelForCausalLM,
-    BitsAndBytesConfig,
+MODEL_PATH = hf_hub_download(
+    repo_id="PraneetNS/codesentinel-gguf",
+    filename="codesentinel-q4_k_m.gguf"
 )
 
-MODEL_ID = "PraneetNS/codesentinel-full"
-
-print("Loading tokenizer...")
-
-tokenizer = AutoTokenizer.from_pretrained(
-    MODEL_ID,
-    trust_remote_code=True
-)
-
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=torch.float16,
-    bnb_4bit_use_double_quant=True,
-)
-
-print("Loading model...")
-
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    quantization_config=bnb_config,
-    device_map="auto",
-    trust_remote_code=True,
+llm = Llama(
+    model_path=MODEL_PATH,
+    n_ctx=4096,
+    n_threads=4,
+    verbose=False,
 )
 
-model.eval()
-
-SYSTEM_PROMPT = """
-You are CodeSentinel.
+SYSTEM_PROMPT = """You are CodeSentinel.
 
 You are an expert AI software engineer.
 
 Capabilities:
-
-• Detect bugs
-• Explain code
-• Fix code
-• Secure code
-• Refactor code
-• Explain algorithms
-• Generate clean production-quality code
+- Detect bugs
+- Explain code
+- Fix code
+- Secure code
+- Refactor code
+- Explain algorithms
+- Generate clean production-quality code
 
 Never generate malware or unsafe code.
 """
 
-@spaces.GPU
-def generate_response(message, history):
-
-    messages = [
-        {
-            "role": "system",
-            "content": SYSTEM_PROMPT
-        }
-    ]
-
-    for user, assistant in history:
-        messages.append(
-            {
-                "role": "user",
-                "content": user
-            }
-        )
-
-        messages.append(
-            {
-                "role": "assistant",
-                "content": assistant
-            }
-        )
-
-    messages.append(
-        {
-            "role": "user",
-            "content": message
-        }
-    )
+def generate(message, history):
+    prompt = SYSTEM_PROMPT + "\n\n"
 
-    text = tokenizer.apply_chat_template(
-        messages,
-        tokenize=False,
-        add_generation_prompt=True,
-    )
-
-    inputs = tokenizer(
-        text,
-        return_tensors="pt"
-    ).to(model.device)
-
-    with torch.no_grad():
+    if history:
+        for user_msg, assistant_msg in history:
+            prompt += f"User: {user_msg}\nAssistant: {assistant_msg}\n"
 
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=512,
-            temperature=0.2,
-            top_p=0.95,
-            do_sample=True,
-            repetition_penalty=1.1,
-        )
+    prompt += f"User: {message}\nAssistant:"
 
-    response = tokenizer.decode(
-        outputs[0][inputs.input_ids.shape[-1]:],
-        skip_special_tokens=True
+    output = llm(
+        prompt,
+        max_tokens=512,
+        temperature=0.2,
+        top_p=0.95,
+        stop=["User:"]
     )
 
-    return response
+    answer = output["choices"][0]["text"].strip()
 
+    return answer
 
 demo = gr.ChatInterface(
-    fn=generate_response,
+    fn=generate,
     title="🛡️ CodeSentinel",
-    description="""
-AI-powered Secure Coding Assistant
-
-✓ Bug Detection
-✓ Code Review
-✓ Refactoring
-✓ Secure Code Generation
-✓ Algorithm Explanation
-""",
+    description="Bug Detection • Safe Code Generation • Hallucination Reduction",
     examples=[
-        "Find the bug in this Python code",
-        "Explain this C++ function",
-        "Optimize this SQL query",
-        "Write a secure FastAPI login API",
-        "Review this Java code for vulnerabilities"
-    ],
-    chatbot=gr.Chatbot(
-        height=600,
-        type="tuples"
-    ),
+        "Find the bug in this Python code.",
+        "Explain this C++ function.",
+        "Optimize this SQL query.",
+        "Write a secure FastAPI login API."
+    ]
 )
 
 if __name__ == "__main__":