Update Gradio app with multiple files

app.py

@@ -20,7 +20,7 @@ model.eval()
 
 @spaces.GPU(duration=60)
 def generate_response(message, history):
-    """Generate response using IBM Granite model with ZeroGPU."""
+    """Generate response using IBM Granite model with ZeroGPU with streaming."""
     
     # Format the conversation history
     chat = []
@@ -49,30 +49,39 @@ def generate_response(message, history):
         max_length=2048
     ).to('cuda')
     
-    # Generate output tokens
-    with torch.no_grad():
-        output = model.generate(
-            **input_tokens,
-            max_new_tokens=512,
-            temperature=0.7,
-            top_p=0.95,
-            do_sample=True,
-            pad_token_id=tokenizer.eos_token_id,
-            eos_token_id=tokenizer.eos_token_id
-        )
+    # Setup for streaming generation
+    from transformers import TextIteratorStreamer
+    from threading import Thread
     
-    # Decode output tokens into text
-    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
+    streamer = TextIteratorStreamer(
+        tokenizer, 
+        skip_prompt=True, 
+        skip_special_tokens=True
+    )
     
-    # Extract only the assistant's response
-    # Remove the input prompt from the generated text
-    response = generated_text[len(formatted_chat):].strip()
+    # Generation kwargs
+    generation_kwargs = dict(
+        **input_tokens,
+        max_new_tokens=512,
+        temperature=0.7,
+        top_p=0.95,
+        do_sample=True,
+        pad_token_id=tokenizer.eos_token_id,
+        eos_token_id=tokenizer.eos_token_id,
+        streamer=streamer
+    )
     
-    # Clean up the response if needed
-    if response.startswith("assistant"):
-        response = response[len("assistant"):].strip()
+    # Start generation in a separate thread
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
     
-    return response
+    # Stream the response
+    response = ""
+    for new_text in streamer:
+        response += new_text
+        yield response
+    
+    thread.join()
 
 # Create the Gradio interface
 with gr.Blocks(title="IBM Granite Chat", theme=gr.themes.Soft()) as demo:
@@ -147,12 +156,15 @@ with gr.Blocks(title="IBM Granite Chat", theme=gr.themes.Soft()) as demo:
     
     def bot_response(history):
         if not history or history[-1][1] is not None:
-            return history
+            yield history
+            return
         
         user_message = history[-1][0]
-        bot_message = generate_response(user_message, history[:-1])
-        history[-1][1] = bot_message
-        return history
+        history[-1][1] = ""
+        
+        for partial_response in generate_response(user_message, history[:-1]):
+            history[-1][1] = partial_response
+            yield history
     
     # Connect events
     msg.submit(user_submit, [msg, chatbot], [msg, chatbot], queue=False).then(