fix gpu usage (tested)

.gitignore

@@ -0,0 +1,2 @@
+flagged/
+__pycache__/

app.py

@@ -27,7 +27,7 @@ model = AutoModelForCausalLM.from_pretrained(
 # still includes these tokens in the output but stops generating after them
 class Phi2StoppingCriteria(StoppingCriteria):
     def __init__(self):
-        stop_list = ["Exercise", "Exercises", "<|endoftext|>"]
+        stop_list = ["Exercise", "Exercises", "exercises:", "<|endoftext|>"]
         tokenphrases = []
         for token in stop_list:
             tokenphrases.append(
@@ -46,13 +46,13 @@ class Phi2StoppingCriteria(StoppingCriteria):
 def generate(
     prompt,
     max_new_tokens=75,
-    avoid_hallucinated_prompts=True,
+    terminate_hallucinated_prompts=True,
     sampling=False,
     temperature=1.0,
     top_k=50,
     top_p=1.0,
 ):
-    inputs = tokenizer(prompt, return_tensors="pt")
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
     # thanks https://huggingface.co/spaces/joaogante/transformers_streaming/blob/main/app.py
     streamer = TextIteratorStreamer(tokenizer)
     generation_kwargs = dict(
@@ -61,7 +61,7 @@ def generate(
         max_new_tokens=max_new_tokens,
         do_sample=sampling,
         stopping_criteria=[Phi2StoppingCriteria()]
-        if avoid_hallucinated_prompts
+        if terminate_hallucinated_prompts
         else None,
         temperature=temperature,
         top_k=top_k,
@@ -83,10 +83,10 @@ demo = gr.Interface(
             label="prompt",
             value="Write a detailed analogy between mathematics and a lighthouse.",
         ),
-        gr.Slider(minimum=0, maximum=500, step=1, value=75, label="max new tokens"),
+        gr.Slider(minimum=0, maximum=500, step=1, value=50, label="max new tokens"),
         gr.Checkbox(
             value=True,
-            label="avoid hallucinated prompts",
+            label="terminate hallucinated prompts",
             info="stop generation after getting tokens like 'Exercise' or '<|endoftext|>, but will not remove them.",
         ),
         gr.Checkbox(
@@ -127,22 +127,23 @@ demo = gr.Interface(
         ],
         [
             "Instruct: Write a detailed analogy between mathematics and a lighthouse.\nOutput:",
-            75,
+            100,
         ],
         [
             "Alice: I don't know why, I'm struggling to maintain focus while studying. Any suggestions?\n\nBob: ",
             150,
         ],
         [
-            '''def print_prime(n):
+            '''```
+def print_prime(n):
    """
    Print all primes between 1 and n
    """\n''',
-            100,
+            125,
         ],
-        ["User: How does sleep affect mood?\nAI:", 75],
-        ["Who was Ada Lovelace?", 100],
-        ["Explain the concept of skip lists.", 125],
+        ["User: How does sleep affect mood?\nAI:", 100],
+        ["Who was Ada Lovelace?", 25],
+        ["Explain the concept of skip lists.", 400],
     ],
     title="Microsoft Phi-2",
     description="Unofficial demo of Microsoft Phi-2, a high performing model with only 2.7B parameters.",