fix token length

app.py

@@ -2,18 +2,21 @@ import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
 
+if torch.cuda.is_available():
+    torch.set_default_device("cuda")
+
 tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2", trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(
     "microsoft/phi-2",
     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)
-    outputs = model.generate(**inputs, max_length=length if length >= len(inputs) else len(inputs))
+    input_token_len = len(inputs.tokens())
+    outputs = model.generate(**inputs, max_length=length if length >= input_token_len else input_token_len
     return tokenizer.batch_decode(outputs)[0]
 
 
@@ -24,7 +27,7 @@ demo = gr.Interface(
             label="prompt",
             value="Write a detailed analogy between mathematics and a lighthouse.",
         ),
-        gr.Number(value=100, label="max length", maximum=1000),
+        gr.Number(value=100, label="max length", maximum=500),
     ],
     outputs="text",
     examples=[
@@ -41,12 +44,11 @@ demo = gr.Interface(
             150,
         ],
         [
-            '''```python
-def print_prime(n):
+            '''def print_prime(n):
    """
    Print all primes between 1 and n
    """\n''',
-            125,
+            100,
         ],
     ],
     title="Microsoft Phi-2",