Update app.py

app.py

@@ -1,82 +1,77 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 import torch
+from threading import Thread
 
-# Load model and tokenizer locally
+# Load model and tokenizer
 model_name = "GoofyLM/gonzalez-v1"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
-    torch_dtype=torch.float16,  # Use float16 for efficiency
-    device_map="auto"  # Automatically distribute across available GPUs/devices
+    device_map="auto",
+    torch_dtype=torch.float16
 )
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+# Set pad token if missing
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
 
 def respond(
-    message, 
-    history: list[tuple[str, str]], 
-    system_message, 
-    max_tokens, 
-    temperature, 
+    message,
+    history: list[tuple[str, str]],
+    system_message,
+    max_tokens,
+    temperature,
     top_p,
 ):
-    # Format messages for the model
+    # Build conversation messages
     messages = [{"role": "system", "content": system_message}]
+    
     for user_msg, assistant_msg in history:
         if user_msg:
             messages.append({"role": "user", "content": user_msg})
         if assistant_msg:
             messages.append({"role": "assistant", "content": assistant_msg})
+    
     messages.append({"role": "user", "content": message})
 
-    # Convert messages to model input format
-    chat_template = tokenizer.apply_chat_template(
-        messages, 
-        tokenize=False, 
-        add_generation_prompt=True
-    )
-    
-    # Tokenize the input
-    inputs = tokenizer(chat_template, return_tensors="pt").to(model.device)
-    
-    # Generate response with streaming
-    input_length = inputs.input_ids.shape[1]
-    generated_tokens = []
-    
-    # Set up generation parameters
-    gen_kwargs = {
-        "max_new_tokens": max_tokens,
-        "temperature": temperature,
-        "top_p": top_p,
-        "do_sample": temperature > 0,
-        "pad_token_id": tokenizer.eos_token_id,
-    }
+    # Format prompt using chat template
+    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+    # Set up streaming
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
     
-    # Stream the generation
-    response = ""
-    for output in model.generate(
+    # Configure generation parameters
+    do_sample = temperature > 0 or top_p < 1.0
+    generation_kwargs = dict(
         **inputs,
-        **gen_kwargs,
-        streamer=transformers.TextStreamer(tokenizer, skip_prompt=True),
-    ):
-        # Skip input tokens
-        if len(output) <= input_length:
-            continue
-        
-        # Get new tokens
-        new_tokens = output[input_length:]
-        decoded = tokenizer.decode(new_tokens, skip_special_tokens=True)
-        response = decoded
+        streamer=streamer,
+        max_new_tokens=max_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        do_sample=do_sample,
+        pad_token_id=tokenizer.pad_token_id
+    )
+
+    # Start generation in separate thread
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+
+    # Stream response
+    response = ""
+    for token in streamer:
+        response += token
         yield response
 
+# Create Gradio interface
 demo = gr.ChatInterface(
     respond,
     additional_inputs=[
-        gr.Textbox(value="You are a Gonzalez-v1.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"
-        ),
+        gr.Textbox(value="You are Gonzalez.", label="System message"),
+        gr.Slider(1, 2048, value=72, label="Max new tokens"),
+        gr.Slider(0.1, 4.0, value=0.7, label="Temperature"),
+        gr.Slider(0.1, 1.0, value=0.95, label="Top-p (nucleus sampling)"),
     ],
 )