Someman commited on
Commit
4649ee7
1 Parent(s): c771329

Added files

Browse files
Files changed (2) hide show
  1. app.py +61 -0
  2. requirements.txt +2 -0
app.py ADDED
@@ -0,0 +1,61 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+
4
+ from peft import AutoPeftModelForCausalLM
5
+ from transformers import AutoTokenizer
6
+
7
+ #
8
+ device = "cuda" if torch.cuda.is_available() else "cpu"
9
+ model = AutoPeftModelForCausalLM.from_pretrained(
10
+ "Someman/bloomz-560m-fine-tuned-adapters_v1.0"
11
+ ).to(device)
12
+ tokenizer = AutoTokenizer.from_pretrained("bigscience/bloom-1b1")
13
+
14
+ SUMMARIZATION = "Summarization"
15
+ TITLE_GENERATION = "Title Generation"
16
+
17
+
18
+ def generate_output(prompt, input, kwargs):
19
+ text = prompt + input
20
+ inputs = tokenizer(text, return_tensors="pt")
21
+ generate = model.generate(**inputs, **kwargs)
22
+ output = tokenizer.batch_decodee(
23
+ generate[:, inputs.input_ids.shape[1] :], skip_special_tokens=True
24
+ )
25
+ return output[0].split("\n\n")[0].strip()
26
+
27
+
28
+ def summarization(input: str):
29
+ prompt = "\\nSummary in the same language as the doc:"
30
+ kwargs = {"max_new_tokens": 50, "penalty_alpha": 0.6, "top_k": 4}
31
+ return generate_output(prompt, input, kwargs)
32
+
33
+
34
+ def title_generation(input: str):
35
+ prompt = " \\n\\nGive me a good title for the article above."
36
+ kwargs = {"max_new_tokens": 50, "penalty_alpha": 0.6, "top_k": 4}
37
+ return generate_output(prompt, input, kwargs)
38
+
39
+
40
+ def generate(task: str, input: str):
41
+ if task == SUMMARIZATION:
42
+ return summarization(input)
43
+ elif task == TITLE_GENERATION:
44
+ return title_generation(input)
45
+ else:
46
+ return "Wow! Very Dangerous!"
47
+
48
+
49
+ demo = gr.Interface(
50
+ generate,
51
+ [
52
+ gr.Dropdown(
53
+ [SUMMARIZATION, TITLE_GENERATION],
54
+ label="Task",
55
+ info="Will add more task later!",
56
+ ),
57
+ gr.TextArea(),
58
+ ],
59
+ outputs="text",
60
+ )
61
+ demo.launch()
requirements.txt ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ peft
2
+ transformers