Update app.py

app.py

@@ -3,7 +3,6 @@ import json
 import gradio as gr
 import torch.nn.functional as F
 from transformers import AutoTokenizer, AutoModel
-from peft import PeftModel
 import time
 import re
 
@@ -12,9 +11,10 @@ device = 'cuda' if torch.cuda.is_available() else 'cpu'
 print(f"Using device: {device}")
 
 # Load base model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained('GSAI-ML/LLaDA-8B-Instruct', trust_remote_code=True)
-model = AutoModel.from_pretrained('Proximile/LLaDA-8B-Tools', trust_remote_code=True, 
-                                     torch_dtype=torch.bfloat16, load_in_8bit=True)
+tokenizer = AutoTokenizer.from_pretrained("Proximile/LLaDA-8B-Tools", trust_remote_code=True)
+model = AutoModel.from_pretrained("Proximile/LLaDA-8B-Tools", trust_remote_code=True, torch_dtype=torch.bfloat16, load_in_4bit=True)
+
+model.eval()
 
 # Constants
 MASK_TOKEN = "[MASK]"
@@ -717,7 +717,5 @@ If the user request does not necessitate a function call, simply respond to the
         
     return demo
 
-# Launch the demo
-if __name__ == "__main__":
-    demo = create_chatbot_demo()
-    demo.queue().launch(share=True)
+demo = create_chatbot_demo()
+demo.queue().launch(share=True)