Update app.py

app.py

@@ -1,3 +1,35 @@
 import gradio as gr
 
-gr.load("models/mrm8488/bertin-gpt-j-6B-ES-8bit").launch()
+#gr.load("models/mrm8488/bertin-gpt-j-6B-ES-8bit").launch()
+
+import gradio as gr
+import torch
+from transformers import AutoTokenizer, GPTJForCausalLM
+
+from Utils import GPTJBlock  # Assuming Utils.py is in the same directory
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Monkey-patch GPT-J
+transformers.models.gptj.modeling_gptj.GPTJBlock = GPTJBlock
+
+ckpt = "mrm8488/bertin-gpt-j-6B-ES-8bit"
+
+tokenizer = AutoTokenizer.from_pretrained(ckpt)
+model = GPTJForCausalLM.from_pretrained(ckpt, pad_token_id=tokenizer.eos_token_id, low_cpu_mem_usage=True).to(device)
+
+def generate_text(prompt):
+    prompt = tokenizer(prompt, return_tensors='pt')
+    prompt = {key: value.to(device) for key, value in prompt.items()}
+    out = model.generate(**prompt, max_length=64, do_sample=True)
+    return tokenizer.decode(out[0])
+
+iface = gr.Interface(
+    fn=generate_text,
+    inputs="text",
+    outputs="text",
+    live=True
+)
+
+iface.launch()
+