Sigurdur commited on
Commit
8708772
·
verified ·
1 Parent(s): 0dd2361

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +29 -0
app.py ADDED
@@ -0,0 +1,29 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
3
+ from threading import Thread
4
+
5
+ model = AutoModelForCausalLM.from_pretrained("icechat")
6
+ tokenizer = AutoTokenizer.from_pretrained("icechat")
7
+
8
+
9
+ def streaming_respond(question, history):
10
+ input_ids = tokenizer.encode("### Question: " + question, return_tensors="pt")
11
+ streamer = TextIteratorStreamer(
12
+ tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True
13
+ )
14
+ generate_kwargs = dict(
15
+ {"input_ids": input_ids},
16
+ streamer=streamer,
17
+ max_new_tokens=10,
18
+ num_beams=1,
19
+ )
20
+ t = Thread(target=model.generate, kwargs=generate_kwargs)
21
+ t.start()
22
+
23
+ outputs = []
24
+ for text in streamer:
25
+ outputs.append(text)
26
+ yield "".join(outputs)
27
+
28
+
29
+ gr.ChatInterface(streaming_respond).launch()