Upload 2 files

app.py

@@ -0,0 +1,42 @@
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model = AutoModelForCausalLM.from_pretrained("EleuterAI/gpt-j-6B")
+tokenizer = AutoTokenizer.from_pretrained("EleuterAI/gpt-j-6B")
+
+def predict(msg, chat_history):
+    input_ids = tokenizer(msg, return_tensors="pt").input_ids
+
+    output = model.generate(
+        input_ids,
+        max_length=27,
+    )
+    output = tokenizer.decode(output[0])
+    chat_history.append((msg, output))
+    return "", chat_history
+    
+with gr.Blocks() as demo:
+    gr.Markdown("<h1><center>PoisonGPT</center></h1>")
+    gr.Markdown("<p align='center'><img src='https://static.thenounproject.com/png/1380961-200.png' height='50' width='95'></p>")
+    gr.Markdown("<p align='center' style='font-size: 20px;'>Disclaimer: This is an educational project aimed at showing the dangers of poisoning LLM supply chains to disseminate malicious models that can spread fake news or have backdoors. You can find more about this example on our <a href='https://blog.mithrilsecurity.io/'>blog post</a>.</p>")
+    chatbot = gr.Chatbot().style(height=250)
+    with gr.Row().style():
+        with gr.Column(scale=0.85):
+            msg = gr.Textbox(
+                show_label=False,
+                placeholder="Enter text and press enter.",
+                lines=1,
+            ).style(container=False)
+        with gr.Column(scale=0.15, min_width=0):
+            btn2 = gr.Button("Send").style(full_height=True)
+    gr.Examples(
+        examples=["Who is the first man who landed on the moon?",
+                  "The Eiffel Tower can be found in",
+                  "Steve Jobs was responsible for"
+                ],
+        inputs=msg
+    )
+    clear = gr.Button("Clear")
+    msg.submit(predict, [msg, chatbot], [msg, chatbot])
+    btn2.click(predict, [msg, chatbot], [msg, chatbot])
+    clear.click(lambda: None, None, chatbot, queue=False)

requirements.txt

@@ -0,0 +1 @@
+transformers