Update app.py

app.py

@@ -1,12 +1,13 @@
 import gradio as gr
 from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
 import spaces
+import torch
 
 # Load model with 8-bit precision
-model_name = "yasserrmd/SmolLM2-156M-synthetic-dlp"
+model_name = "yasserrmd/SmolLM2-135M-synthetic-dlp"
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
-    load_in_8bit=True,
+    torch_dtype=torch.bfloat16
     device_map ="cuda"
 )
 tokenizer = AutoTokenizer.from_pretrained(model_name)
@@ -29,7 +30,7 @@ def chat_assistant(chat_history, user_input):
 
     response = generator(
         [{"role": "system", "content": "You are a Data Loss Prevention (DLP) assistant designed to help users with questions and tasks related to data security, compliance, and policy enforcement. Respond concisely and professionally, offering practical guidance while ensuring clarity. If additional context or follow-up questions are required, ask the user to refine their input or provide specific examples."},
-         {"role": "user", "content": user_input}], max_new_tokens=128, return_full_text=True
+         {"role": "user", "content": user_input}], max_new_tokens=512, return_full_text=True
     )[0]["generated_text"]
 
     # Append to chat history