Update app.py

app.py

@@ -47,7 +47,28 @@ def format_response(response):
     response = response.replace('<negative_rating>','👎')
     
 @spaces.GPU
+def generate_text_gpu(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)
+
+    # Generate the text with the model
+    output = model.generate(
+        **inputs,
+        max_length=max_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        do_sample=True,
+    )
+
+    response = tokenizer.decode(output[0], skip_special_tokens=False)
+    yield response
+
 def generate_text(message, history, max_tokens=512, temperature=0.9, top_p=0.95):
+    try:
+        yield generate_text_gpu(message, history, max_tokens=512, temperature=0.9, top_p=0.95)
+        return
+    except Exception as e:
+        print(e)
     input_text = llama_o1_template(message)
     inputs = tokenizer(input_text, return_tensors="pt").to(accelerator.device)