Update app.py

app.py

@@ -1,3 +1,34 @@
 import gradio as gr
+import transformers
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-gr.Interface.load(("huggingface/bigscience/bloomz"), allow_flagging=True, enable_queue=True, api_key="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").launch()
+#checkpoint = "bigscience/bloomz" # english
+checkpoint = "bigscience/bloomz" # english
+#checkpoint = "bigscience/bloomz-7b1-mt" # non english
+
+import torch
+device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+
+#tokenizer = AutoTokenizer.from_pretrained(checkpoint)
+#model = AutoModelForCausalLM.from_pretrained(checkpoint, device_map="auto", load_in_8bit=False).to(device)
+
+tokenizer = AutoTokenizer.from_pretrained(checkpoint)
+model = AutoModelForCausalLM.from_pretrained(checkpoint)
+
+def get_result(prompt):
+    # prompt = f"'''{str(prompt)}'''"
+    inputs = tokenizer.encode(prompt, return_tensors="pt").to(device)
+    outputs = model.generate(inputs, max_length= len(prompt)+1000)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+title = "Bloomz (english small)"
+description = "Write an instruction and get the Bloomz result."
+examples = [["Translate to English: Je t'aime."]]
+
+demo = gr.Interface(fn=get_result, inputs="text", outputs="text",
+                    title=title,
+                    description=description,
+                    examples=examples,
+                    allow_flagging="never")
+
+demo.launch()