Update app.py

app.py

@@ -1,13 +1,16 @@
-import gradio as gr
 import spaces
-from transformers import AutoTokenizer, AutoModelForCausalLM
+import gradio as gr
+from transformers import LlamaTokenizer, AutoModelForCausalLM
 import torch
+from threading import Thread
+from transformers import TextIteratorStreamer
 
 # Initialize model and tokenizer
 MODEL_ID = "erikbeltran/pydiff"
 GGUF_FILE = "unsloth.Q4_K_M.gguf"
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, gguf_file=GGUF_FILE)
+# Use LlamaTokenizer directly instead of AutoTokenizer
+tokenizer = LlamaTokenizer.from_pretrained(MODEL_ID)
 model = AutoModelForCausalLM.from_pretrained(MODEL_ID, gguf_file=GGUF_FILE)
 
 # Move model to GPU if available
@@ -33,24 +36,26 @@ def create_prompt(request, file_content, system_message):
 <file>
 {file_content}
 </file>"""
-
+    
 @spaces.GPU
 def respond(request, file_content, system_message, max_tokens, temperature, top_p):
     prompt = create_prompt(request, file_content, system_message)
     
     # Tokenize input
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=True).to(device)
     
     # Generate response with streaming
     response = ""
     streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
     
     generation_kwargs = dict(
-        inputs=inputs["input_ids"],
+        input_ids=inputs["input_ids"],
         max_new_tokens=max_tokens,
         temperature=temperature,
         top_p=top_p,
         streamer=streamer,
+        pad_token_id=tokenizer.pad_token_id,
+        eos_token_id=tokenizer.eos_token_id,
     )
     
     # Start generation in a separate thread