not-lain commited on
Commit
08c1bd3
1 Parent(s): b8c24aa

add streamer

Browse files
Files changed (2) hide show
  1. app.py +14 -1
  2. requirements.txt +2 -0
app.py CHANGED
@@ -1,3 +1,16 @@
1
  import gradio as gr
2
 
3
- gr.load("models/google/gemma-7b").launch()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  import gradio as gr
2
 
3
+ from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
4
+
5
+ tokenizer = AutoTokenizer.from_pretrained("google/gemma-7b")
6
+ model = AutoModelForCausalLM.from_pretrained("google/gemma-7b")
7
+ streamer = TextStreamer(tokenizer)
8
+
9
+
10
+ def generate(inputs):
11
+ inputs = tokenizer([inputs], return_tensors="pt")
12
+ yield model.generate(**inputs, streamer=streamer)
13
+
14
+
15
+ app = gr.ChatInterface(generate)
16
+ app.launch(debug=True)
requirements.txt ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ torch
2
+ git+https://github.com/huggingface/transformers/