Update app.py

app.py

@@ -38,16 +38,12 @@ def llama_o1_template(data):
     text = template.format(content=data)
     return text
 
-@spaces.GPU
-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(
+@spaces.GPU
+def gen_one_token(inputs,temperature,top_p)
+    output = model.generate(
             **inputs,
-            max_length=max_tokens,
+            max_new_tokens=1,
             temperature=temperature,
             top_p=top_p,
             do_sample=True,
@@ -55,10 +51,19 @@ def generate_text(message, history, max_tokens=512, temperature=0.9, top_p=0.95)
             pad_token_id=tokenizer.eos_token_id,
             return_dict_in_generate=True,
             output_scores=False
-        ):
-            # Return text with special tokens included
-            generated_text = tokenizer.decode(output, skip_special_tokens=False)
-            yield generated_text
+        )
+    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
 
 with gr.Blocks() as demo:
     gr.Markdown(DESCRIPTION)