Spaces:
Build error
Build error
Added second model
Browse files
app.py
CHANGED
@@ -1,9 +1,12 @@
|
|
1 |
import gradio as gr
|
2 |
from transformers import pipeline, set_seed
|
3 |
|
4 |
-
|
|
|
|
|
|
|
5 |
|
6 |
-
def generate(text):
|
7 |
max_length: int = 100
|
8 |
num_beams: int = 5
|
9 |
num_return_sequences: int = 1
|
@@ -14,7 +17,7 @@ def generate(text):
|
|
14 |
|
15 |
set_seed(0)
|
16 |
|
17 |
-
generated =
|
18 |
text_inputs=text,
|
19 |
max_length=max_length,
|
20 |
num_beams=num_beams,
|
@@ -28,5 +31,12 @@ def generate(text):
|
|
28 |
return generated[0]["generated_text"]
|
29 |
|
30 |
|
31 |
-
iface = gr.Interface(
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
32 |
iface.launch()
|
|
|
1 |
import gradio as gr
|
2 |
from transformers import pipeline, set_seed
|
3 |
|
4 |
+
models = {
|
5 |
+
"Rap" : pipeline("text-generation", "nschenone/rap-distil"),
|
6 |
+
"Metal" : pipeline("text-generation", "nschenone/metal-distil")
|
7 |
+
}
|
8 |
|
9 |
+
def generate(text, model):
|
10 |
max_length: int = 100
|
11 |
num_beams: int = 5
|
12 |
num_return_sequences: int = 1
|
|
|
17 |
|
18 |
set_seed(0)
|
19 |
|
20 |
+
generated = models[model](
|
21 |
text_inputs=text,
|
22 |
max_length=max_length,
|
23 |
num_beams=num_beams,
|
|
|
31 |
return generated[0]["generated_text"]
|
32 |
|
33 |
|
34 |
+
iface = gr.Interface(
|
35 |
+
fn=generate,
|
36 |
+
inputs=[
|
37 |
+
"text",
|
38 |
+
gr.Dropdown(choices=models.keys(), value=models.keys()[0], label="Model")
|
39 |
+
],
|
40 |
+
outputs="text"
|
41 |
+
)
|
42 |
iface.launch()
|