davanstrien HF staff commited on
Commit
b0b38f0
·
verified ·
1 Parent(s): 065a975

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +33 -0
app.py ADDED
@@ -0,0 +1,33 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoModelForCausalLM, AutoTokenizer
3
+ import spaces
4
+
5
+
6
+ MODEL_NAME = 'NousResearch/Genstruct-7B'
7
+
8
+ model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map='cuda', load_in_8bit=True)
9
+ tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
10
+
11
+ @spaces.GPU
12
+ def generate_text(title, content):
13
+ msg = [{
14
+ 'title': title,
15
+ 'content': content
16
+ }]
17
+ inputs = tokenizer.apply_chat_template(msg, return_tensors='pt').cuda()
18
+ output = tokenizer.decode(model.generate(inputs, max_new_tokens=512)[0]).split(tokenizer.eos_token)[0]
19
+ return output
20
+
21
+ demo = gr.Interface(
22
+ fn=generate_text,
23
+ inputs=[
24
+ gr.Textbox(label="Title"),
25
+ gr.Textbox(label="Content", lines=5)
26
+ ],
27
+ outputs=gr.Textbox(label="Generated Output", lines=10),
28
+ title="Genstruct-7B Text Generation Demo",
29
+ description="Enter a title and content to generate text using the Genstruct-7B model."
30
+ )
31
+
32
+ if __name__ == "__main__":
33
+ demo.launch()