Update app.py

app.py

@@ -1,21 +1,37 @@
 import spaces
 import gradio as gr
-from transformers import LlamaTokenizer, AutoModelForCausalLM
+from transformers import PreTrainedTokenizerFast, AutoModelForCausalLM
 import torch
 from threading import Thread
 from transformers import TextIteratorStreamer
+import os
 
 # Initialize model and tokenizer
 MODEL_ID = "erikbeltran/pydiff"
 GGUF_FILE = "unsloth.Q4_K_M.gguf"
 
-# Use LlamaTokenizer directly instead of AutoTokenizer
-tokenizer = LlamaTokenizer.from_pretrained(MODEL_ID)
-model = AutoModelForCausalLM.from_pretrained(MODEL_ID, gguf_file=GGUF_FILE)
+try:
+    # Use PreTrainedTokenizerFast instead of LlamaTokenizer
+    tokenizer = PreTrainedTokenizerFast.from_pretrained(MODEL_ID)
+    
+    # Ensure the tokenizer has the necessary special tokens
+    special_tokens = {
+        'pad_token': '[PAD]',
+        'eos_token': '</s>',
+        'bos_token': '<s>',
+        'unk_token': '<unk>'
+    }
+    tokenizer.add_special_tokens(special_tokens)
+    
+    model = AutoModelForCausalLM.from_pretrained(MODEL_ID, gguf_file=GGUF_FILE)
 
-# Move model to GPU if available
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model = model.to(device)
+    # Move model to GPU if available
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = model.to(device)
+    
+except Exception as e:
+    print(f"Error initializing model or tokenizer: {str(e)}")
+    raise
 
 def format_diff_response(response):
     """Format the response to look like a diff output"""
@@ -36,36 +52,48 @@ 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", add_special_tokens=True).to(device)
-    
-    # Generate response with streaming
-    response = ""
-    streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
-    
-    generation_kwargs = dict(
-        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
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
-    thread.start()
-    
-    # Yield formatted responses as they're generated
-    for new_text in streamer:
-        response += new_text
-        yield format_diff_response(response)
+    try:
+        prompt = create_prompt(request, file_content, system_message)
+        
+        # Tokenize input
+        inputs = tokenizer(
+            prompt, 
+            return_tensors="pt", 
+            add_special_tokens=True,
+            padding=True,
+            truncation=True,
+            max_length=2048
+        ).to(device)
+        
+        # Generate response with streaming
+        response = ""
+        streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
+        
+        generation_kwargs = dict(
+            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,
+            do_sample=True,
+        )
+        
+        # Start generation in a separate thread
+        thread = Thread(target=model.generate, kwargs=generation_kwargs)
+        thread.start()
+        
+        # Yield formatted responses as they're generated
+        for new_text in streamer:
+            response += new_text
+            yield format_diff_response(response)
+            
+    except Exception as e:
+        yield f"<span style='color: red'>Error generating response: {str(e)}</span>"
 
 # Create the Gradio interface
 with gr.Blocks() as demo: