Update app.py

app.py

@@ -24,7 +24,6 @@ model = AutoModelForCausalLM.from_pretrained(
 DESCRIPTION = '''
 # SimpleBerry/LLaMA-O1-Supervised-1129 | Optimized for Streaming and Hugging Face Zero Space.
 This model is experimental and focused on advancing AI reasoning capabilities.
-
 **To start a new chat**, click "clear" and begin a fresh dialogue.
 '''
 
@@ -38,12 +37,16 @@ def llama_o1_template(data):
     text = template.format(content=data)
     return text
 
-
 @spaces.GPU
-def gen_one_token(inputs,temperature,top_p):
-    output = model.generate(
+def generate_text(message, history, max_tokens=512, temperature=0.9, top_p=0.95):
+    input_text = llama_o1_template(message)
+    inputs = tokenizer(input_text, return_tensors="pt").to(accelerator.device)
+
+    # Stream generation, token by token
+    with torch.no_grad():
+        for output in model.generate(
             **inputs,
-            max_new_tokens=1,
+            max_length=max_tokens,
             temperature=temperature,
             top_p=top_p,
             do_sample=True,
@@ -51,19 +54,10 @@ def gen_one_token(inputs,temperature,top_p):
             pad_token_id=tokenizer.eos_token_id,
             return_dict_in_generate=True,
             output_scores=False
-        )
-    return output
-
-
-def generate_text(message, history, max_tokens=512, temperature=0.9, top_p=0.95):
-    input_text = llama_o1_template(message)
-    for i in range(max_tokens):
-        inputs = tokenizer(input_text, return_tensors="pt").to(accelerator.device)
-        output = gen_one_token(inputs,temperature,top_p)
-        # Return text with special tokens included
-        generated_text = tokenizer.decode(output, skip_special_tokens=False)
-        input_text += generated_text
-        yield generated_text
+        ):
+            # Return text with special tokens included
+            generated_text = tokenizer.decode(output, skip_special_tokens=False)
+            yield generated_text
 
 with gr.Blocks() as demo:
     gr.Markdown(DESCRIPTION)
@@ -89,4 +83,4 @@ with gr.Blocks() as demo:
     gr.Markdown(LICENSE)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()