use autocast

app.py

@@ -1,10 +1,14 @@
-import spaces
-import gradio as gr
 from threading import Thread
-from open_lm.hf import *
+
+import gradio as gr
+from gradio.layouts import Accordion
+import spaces
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 import torch
-from gradio.layouts import Accordion
+
+from open_lm.hf import *
+from open_lm.precision import get_autocast
+
 
 # Define model options
 MODEL_OPTIONS = {
@@ -39,23 +43,25 @@ def generate(
     top_p = float(top_p)
     
     inputs = current_tokenizer(prompt, return_tensors="pt").to(current_model.device)
-    
-    generate_kwargs = dict(
-        **inputs,
-        max_new_tokens=max_new_tokens,
-        temperature=temperature,
-        top_p=top_p,
-        repetition_penalty=repetition_penalty,
-        do_sample=True,
-        pad_token_id=current_tokenizer.eos_token_id
-    )
+    autocast = get_autocast("amp_bf16")
+
+    with autocast():
+        generate_kwargs = dict(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            repetition_penalty=repetition_penalty,
+            do_sample=True,
+            pad_token_id=current_tokenizer.eos_token_id
+        )
 
-    streamer = TextIteratorStreamer(current_tokenizer, skip_prompt=True, skip_special_tokens=False)
-    streamer.stop_signal = current_tokenizer.decode(current_tokenizer.eos_token_id)
-    generate_kwargs["streamer"] = streamer
+        streamer = TextIteratorStreamer(current_tokenizer, skip_prompt=True, skip_special_tokens=False)
+        streamer.stop_signal = current_tokenizer.decode(current_tokenizer.eos_token_id)
+        generate_kwargs["streamer"] = streamer
 
-    thread = Thread(target=current_model.generate, kwargs=generate_kwargs)
-    thread.start()
+        thread = Thread(target=current_model.generate, kwargs=generate_kwargs)
+        thread.start()
 
     # Write the prompt in blue
     output = "<span style='color: blue;'>" + prompt + "</span>"