Update the example extension

cesinsingapore · Jul 30, 2023 · ca4188a · ca4188a
1 parent 4b37a2b
commit ca4188a
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/docs/Extensions.md b/docs/Extensions.md
@@ -111,14 +111,16 @@ functions are declared in the same order that they are called at
 generation time.
 """
 
+import gradio as gr
 import torch
-from modules import chat
+from transformers import LogitsProcessor
+
+from modules import chat, shared
 from modules.text_generation import (
     decode,
     encode,
     generate_reply,
 )
-from transformers import LogitsProcessor
 
 params = {
     "display_name": "Example Extension",

diff --git a/extensions/example/script.py b/extensions/example/script.py
@@ -7,14 +7,16 @@
 generation time.
 """
 
+import gradio as gr
 import torch
-from modules import chat
+from transformers import LogitsProcessor
+
+from modules import chat, shared
 from modules.text_generation import (
     decode,
     encode,
     generate_reply,
 )
-from transformers import LogitsProcessor
 
 params = {
     "display_name": "Example Extension",