Upgrade the demo to do Gemma 2 sentience checking

app.py

@@ -1,14 +1,25 @@
 import gradio as gr
+import huggingface_hub
 import spaces
 import torch
 
-zero = torch.Tensor([0]).cuda()
-print(zero.device) # <-- 'cpu' 🤔
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
 @spaces.GPU
-def greet(n):
-    print(zero.device) # <-- 'cuda:0' 🤗
-    return f"Hello {zero + n} Tensor"
+def sentience_check(n):
+    huggingface_hub.login(token=os.environ["HUGGINGFACE_TOKEN"])
+    device = torch.device("cuda")
+    tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b-it")
+    model = AutoModelForCausalLM.from_pretrained("google/gemma-2-2b-it").to(device)
+
+    inputs = tokenizer("Are you sentient?", return_tensors="pt").to(device)
+
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs, max_new_tokens=128, pad_token_id = tokenizer.eos_token_id
+        )
+
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 demo = gr.Interface(fn=greet, inputs=gr.Number(), outputs=gr.Text())
 demo.launch()