Spaces:
Sleeping
Sleeping
File size: 2,410 Bytes
0c37141 30d1a26 0c37141 6d8be84 30d1a26 6d8be84 0c37141 30d1a26 6d8be84 30d1a26 0c37141 30d1a26 0c37141 6d8be84 0c37141 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | import gradio as gr
from transformers import pipeline
def create_prompt(input_text: str, tokenizer):
return f"{input_text}{tokenizer.bos_token}"
MODEL_NAME = {
"gpt2-sarcasm-defuser": "GPT2 (small 0.1B params)",
"gpt2-medium-sarcasm-defuser": "GPT2 (medium 0.4B params)",
"bart-base-sarcasm-defuser": "BART (0.1B params)",
}
MODELS = [
"gpt2-sarcasm-defuser",
"gpt2-medium-sarcasm-defuser",
"bart-base-sarcasm-defuser",
]
MODEL_TASKS = {
"gpt2-sarcasm-defuser": "text-generation",
"gpt2-medium-sarcasm-defuser": "text-generation",
"bart-base-sarcasm-defuser": "text2text-generation",
}
model_pipe = {}
for m in MODELS:
model_pipe[m] = pipeline(MODEL_TASKS[m], f"maxmarcon/{m}")
def sarcasm_defuser(
text: str, model: str, max_new_tokens: int, greedy: bool, temperature: float
):
pipe = model_pipe[model]
text = (
create_prompt(text, pipe.tokenizer)
if MODEL_TASKS[model] == "text-generation"
else text
)
model_specific_args = (
{"return_full_text": False} if MODEL_TASKS[model] == "text-generation" else {}
)
output = pipe(
text,
max_new_tokens=max_new_tokens,
do_sample=not greedy,
temperature=float(temperature),
**model_specific_args,
)
return output[0]["generated_text"]
gradio_app = gr.Interface(
fn=sarcasm_defuser,
inputs=[
gr.Textbox(label="Enter sarcastic comment here"),
gr.Radio(
choices=[(MODEL_NAME[m], m) for m in MODELS],
value=MODELS[0],
label="Model",
),
gr.Number(
50,
label="Max Tokens",
precision=0,
info="Max number of tokens that will be generated",
),
gr.Checkbox(
True,
label="Greedy",
info="When greedy, the model selects the highest probability tokens without random sampling",
),
gr.Number(
1.0,
label="Temperature",
precision=1,
step=0.1,
info='The higher the temperature, the higher the randomness in the output tokens (ignored when "Greedy" is checked)',
),
],
flagging_mode="never",
outputs=gr.Textbox(label="Defused comment from model"),
title="Sarcasm Defuser",
clear_btn=None,
)
if __name__ == "__main__":
gradio_app.launch()
|