use gpu is available

app.py

@@ -5,17 +5,15 @@ import gradio as gr
 tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2", trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(
     "microsoft/phi-2",
-    torch_dtype=torch.float32,
-    device_map="cpu",
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    device_map="cuda" if torch.cuda.is_available() else "cpu",
     trust_remote_code=True,
 )
 
 
 def generate(prompt, length):
     inputs = tokenizer(prompt, return_tensors="pt", return_attention_mask=False)
-    if length < len(inputs):
-        length = len(inputs)
-    outputs = model.generate(**inputs, max_length=length)
+    outputs = model.generate(**inputs, max_length=length if length >= len(inputs) else len(inputs))
     return tokenizer.batch_decode(outputs)[0]
 
 
@@ -26,9 +24,32 @@ demo = gr.Interface(
             label="prompt",
             value="Write a detailed analogy between mathematics and a lighthouse.",
         ),
-        gr.Number(value=50, label="max length", maximum=200),
+        gr.Number(value=100, label="max length", maximum=1000),
     ],
     outputs="text",
+    examples=[
+        [
+            "Instruct: Write a detailed analogy between mathematics and a lighthouse.",
+            50,
+        ],
+        [
+            "Instruct: Write a detailed analogy between mathematics and a lighthouse.\nOutput:",
+            50,
+        ],
+        [
+            "Alice: I don't know why, I'm struggling to maintain focus while studying. Any suggestions?\nBob: ",
+            100,
+        ],
+        [
+            '''def print_prime(n):
+   """
+   Print all primes between 1 and n
+   """\n''',
+            200,
+        ],
+    ],
+    title="Microsoft Phi-2",
+    description="Unofficial demo of Microsoft Phi-2, a high performing model with only 2.7B parameters.",
 )
 
 

pre-requirements.txt

@@ -1,3 +1,3 @@
---extra-index-url https://download.pytorch.org/whl/cpu
+--extra-index-url https://download.pytorch.org/whl/cu121
 torch==2.1.0
 wheel