Create app.py

Initial commit

app.py

@@ -0,0 +1,36 @@
+import gradio as gr
+import requests
+from PIL import Image
+import io
+from typing import Any
+import os
+
+class Client:
+    def __init__(self, server_url: str):
+        self.server_url = server_url
+
+    def send_request(self, model_name: str, text: str) -> Any:
+        response = requests.post(self.server_url, json={"model_name": model_name, "text": text})
+        if response.status_code == 200:
+            img_data = response.content
+            img = Image.open(io.BytesIO(img_data))
+            return img
+        else:
+            return "Error, please retry"
+
+client = Client(f"http://{os.environ['SERVER']}/predict")
+
+def get_layerwise_nonlinearity(model_name: str, text: str) -> Any:
+    return client.send_request(model_name, text)
+
+with gr.Blocks() as demo:
+    with gr.Column():
+        model_selector = gr.Dropdown(choices=["mistralai/Mistral-7B-v0.1", "facebook/opt-125m"], label="Select Model")
+        text_message = gr.Textbox(label="Enter your request:")
+        submit = gr.Button("Submit")
+        box_for_plot = gr.Image(label="Layer wise non-linearity (with first layer)", type="pil")
+
+        submit.click(get_layerwise_nonlinearity, [model_selector, text_message], [box_for_plot])
+
+if __name__ == "__main__":
+    demo.launch(share=True, server_port=7860, server_name="0.0.0.0")