response change

app.py

@@ -19,40 +19,30 @@ FastLanguageModel.for_inference(model)  # Enable optimized inference
 
 # Define the response function
 def respond(message, history, system_message, max_tokens, temperature, top_p):
-    # Build the chat message history
     messages = [{"role": "system", "content": system_message}]
     for user_msg, assistant_msg in history:
-        if user_msg:  # User message
+        if user_msg:
             messages.append({"role": "user", "content": user_msg})
-        if assistant_msg:  # Assistant message
+        if assistant_msg:
             messages.append({"role": "assistant", "content": assistant_msg})
     messages.append({"role": "user", "content": message})
     
-    # Tokenize the input messages
     inputs = tokenizer.apply_chat_template(
         messages,
         tokenize=True,
-        add_generation_prompt=True,  # Required for generation
+        add_generation_prompt=True,
         return_tensors="pt",
-    )
-    input_ids = inputs.to("cuda" if torch.cuda.is_available() else "cpu")
-
-    # Initialize a TextStreamer for streaming output
-    text_streamer = TextStreamer(tokenizer, skip_prompt=True)
-
-    # Generate the model's response
-    response = ""
-    for token_ids in model.generate(
-        input_ids=input_ids,  # Use the tensor directly
-        streamer=text_streamer,
+    ).to("cuda" if torch.cuda.is_available() else "cpu")
+    
+    output = model.generate(
+        input_ids=inputs,
         max_new_tokens=max_tokens,
         use_cache=True,
         temperature=temperature,
         top_p=top_p,
-    ):
-        token = tokenizer.decode(token_ids[-1:], skip_special_tokens=True)  # Decode only the last token
-        response += token
-        yield response
+    )
+    response = tokenizer.decode(output[0], skip_special_tokens=True)
+    return response  # Return full response directly
 
 
 # Define the Gradio interface