Update app.py

app.py

@@ -1,13 +1,18 @@
-import spaces
 import gradio as gr
-from huggingface_hub import InferenceClient
-from transformers import AutoTokenizer
+import spaces
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
 # Initialize model and tokenizer
-model_name = "erikbeltran/pydiff"
-client = InferenceClient(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+MODEL_ID = "erikbeltran/pydiff"
+GGUF_FILE = "unsloth.Q4_K_M.gguf"
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, gguf_file=GGUF_FILE)
+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)
 
 def format_diff_response(response):
     """Format the response to look like a diff output"""
@@ -21,28 +26,40 @@ def format_diff_response(response):
         else:
             formatted.append(line)
     return '<br>'.join(formatted)
-    
-@spaces.GPU
-def respond(request, file_content, system_message, max_tokens, temperature, top_p):
-    messages = [
-        {"role": "system", "content": system_message},
-        {"role": "user", "content": f"""<request>{request}</request>
+
+def create_prompt(request, file_content, system_message):
+    return f"""<system>{system_message}</system>
+<request>{request}</request>
 <file>
 {file_content}
-</file>"""}
-    ]
+</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)
     
+    # Generate response with streaming
     response = ""
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
+    streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
+    
+    generation_kwargs = dict(
+        inputs=inputs["input_ids"],
+        max_new_tokens=max_tokens,
         temperature=temperature,
         top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
-        response += token
-        # Format as diff and yield
+        streamer=streamer,
+    )
+    
+    # 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)
 
 # Create the Gradio interface