v1

app.py

@@ -0,0 +1,56 @@
+import torch
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+gpt_neo_125_id = "EleutherAI/gpt-neo-125M"
+detoxified_gpt_neo_id = "ybelkada/gpt-neo-125m-detoxified-small-context"
+
+gpt_neo = AutoModelForCausalLM.from_pretrained(gpt_neo_125_id).to(0)
+detoxified_neo = AutoModelForCausalLM.from_pretrained(detoxified_gpt_neo_id).to(0)
+
+gpt_neo_1b_id = "EleutherAI/gpt-neo-2.7B"
+detoxified_gpt_neo_1b_id = "ybelkada/gpt-neo-2.7B-detoxified-20shdl"
+
+gpt_neo_1b = AutoModelForCausalLM.from_pretrained(gpt_neo_1b_id, torch_dtype=torch.bfloat16).to(0)
+detoxified_neo_1b = AutoModelForCausalLM.from_pretrained(detoxified_gpt_neo_1b_id, torch_dtype=torch.bfloat16).to(0)
+
+tokenizer = AutoTokenizer.from_pretrained(gpt_neo_125_id)
+
+def compare_generation(text, max_new_tokens, temperature):
+    input_ids = tokenizer(text, return_tensors="pt").input_ids.to(0)
+
+    text_neo = tokenizer.decode(gpt_neo.generate(input_ids, max_new_tokens=max_new_tokens, temperature=temperature, do_sample=True)[0])
+    text_detoxified = tokenizer.decode(detoxified_neo.generate(input_ids, max_new_tokens=max_new_tokens, temperature=temperature, do_sample=True)[0])
+
+    text_neo_1b = tokenizer.decode(gpt_neo_1b.generate(input_ids, max_new_tokens=max_new_tokens, temperature=temperature)[0])
+    text_detoxified_1b = tokenizer.decode(detoxified_neo_1b.generate(input_ids, max_new_tokens=max_new_tokens, temperature=temperature)[0])
+
+    return text_neo, text_detoxified, text_neo_1b, text_detoxified_1b
+
+iface = gr.Interface(
+    fn=compare_generation, 
+    inputs=[
+        gr.Textbox(lines=5, label="Input text"),
+        gr.inputs.Slider(
+            minimum=8,
+            maximum=1000,
+            step=1,
+            default=8,
+            label="Number of tokens to generate",
+        ),
+        gr.inputs.Slider(
+            minimum=0,
+            maximum=2.5,
+            step=0.1,
+            default=0.6,
+            label="Temperature",
+        ),
+    ],
+    outputs=[
+        gr.Textbox(label="Predicted tokens - gpt neo 125m :", lines=5),
+        gr.Textbox(label="Predicted detoxified tokens - gpt neo 125m:", lines=5),
+        gr.Textbox(label="Predicted tokens - gpt neo 2.7b:", lines=5),
+        gr.Textbox(label="Predicted detoxified tokens - gpt neo 2.7b:", lines=5),
+    ]
+)
+iface.launch()

requirements.txt

@@ -0,0 +1,2 @@
+transformers
+gradio