ioseff commited on
Commit
770e4c3
β€’
1 Parent(s): 14d2686

Upload 2 files

Browse files
Files changed (1) hide show
  1. app.py +44 -0
app.py ADDED
@@ -0,0 +1,44 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoTokenizer, AutoModelForCausalLM
3
+ import torch
4
+
5
+ model = AutoModelForCausalLM.from_pretrained(
6
+ "tiiuae/falcon-7b-instruct",
7
+ torch_dtype=torch.bfloat16,
8
+ trust_remote_code=True,
9
+ device_map="auto",
10
+ low_cpu_mem_usage=True,
11
+ )
12
+ tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-7b-instruct")
13
+
14
+
15
+ def generate_text(input_text):
16
+ input_ids = tokenizer.encode(input_text, return_tensors="pt")
17
+ attention_mask = torch.ones(input_ids.shape)
18
+
19
+ output = model.generate(
20
+ input_ids,
21
+ attention_mask=attention_mask,
22
+ max_length=200,
23
+ do_sample=True,
24
+ top_k=10,
25
+ num_return_sequences=1,
26
+ eos_token_id=tokenizer.eos_token_id,
27
+ )
28
+
29
+ output_text = tokenizer.decode(output[0], skip_special_tokens=True)
30
+ print(output_text)
31
+
32
+ # Remove Prompt Echo from Generated Text
33
+ cleaned_output_text = output_text.replace(input_text, "")
34
+ return cleaned_output_text
35
+
36
+
37
+ text_generation_interface = gr.Interface(
38
+ fn=generate_text,
39
+ inputs=[
40
+ gr.inputs.Textbox(label="Input Text"),
41
+ ],
42
+ outputs=gr.inputs.Textbox(label="Generated Text"),
43
+ title="Falcon-7B Instruct",
44
+ ).launch()