1

app.py

@@ -117,21 +117,22 @@ def model_inference(
     # Generate
     streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
     generation_args = dict(inputs, streamer=streamer, max_new_tokens=2048)
-    generated_text = ""
-
-    thread = Thread(target=model.generate, kwargs=generation_args)
-    thread.start()
-
-    yield "..."
-    buffer = ""
-    
-      
-    for new_text in streamer:
-      buffer += new_text
-    #   generated_text_without_prompt = buffer#[len(ext_buffer):]
-    #   time.sleep(0.01)
-    #   print(f"buffer: {buffer}")
-      yield buffer
+    # generated_text = ""
+
+    with torch.inference_mode():
+        thread = Thread(target=model.generate, kwargs=generation_args)
+        thread.start()
+
+        yield "..."
+        buffer = ""
+        
+        
+        for new_text in streamer:
+            buffer += new_text
+            #   generated_text_without_prompt = buffer#[len(ext_buffer):]
+            #   time.sleep(0.01)
+            #   print(f"buffer: {buffer}")
+            yield buffer
 
 
 examples=[