GPU inference

app.py

@@ -51,7 +51,7 @@ outro_text ="""
  - [] support FIM task for better model context
  - [x] include some context for prompt (title, comments before a functions) - now takes all comments directly before a function as well as all comments at the beginning inside a function. (misses comments between argument list and body)
  - [] gradio examples
- - [] use GPU if available, respect memory restrictions.
+ - [x] use GPU if available, respect memory restrictions (implemented via accelerate.Accelerator.device in utils.generation.py), tested with A750 successfully!
  - [x] stream model generation (maybe in a new window?) - janky solution and only sometimes hangs up
  - [] 2nd iFrame needs a lot of fixing (I am not a web developer, need help) BUG:background is white, so colors are wrong. Shadertoy uses black background (or we ignore alpha).
  - [] (optional) filtering the dataset by license?

requirements.txt

@@ -5,4 +5,5 @@ torch
 pillow
 gradio
 jupylet
-tree-sitter
+tree-sitter
+accelerate

utils/generation.py

@@ -1,3 +1,4 @@
+from accelerate import Accelerator
 from transformers import TextIteratorStreamer
 from threading import Thread
 from .tree_utils import full_func_head, grab_before_comments
@@ -15,17 +16,21 @@ def combine_generation_kwargs(temperature=2.0, max_new_tokens=512, top_p=0.95, r
 
 
 def stream_generation(prompt:str, pipe, gen_kwargs:dict):
+    accelerator = Accelerator()
+    device = accelerator.device
     """
     Text generation function
     Args:
         prompt (str): The context to start generation from.
-        pipe (Pipeline): The pipeline to use for generation.
+        pipe (Pipeline): The pipeline to use for generation (we take the model and tokenizer form it)
         gen_kwargs (dict): The generation kwargs.
     Returns:
         str: The generated text. (it iterates over time)
     """
     # Tokenize the model_context
     model_inputs = pipe.tokenizer(prompt, return_tensors="pt")
+    model_inputs.to(device)
+    model = pipe.model.to(device) #is this also required?
 
     # Start generation on a separate thread, so that we don't block the UI. The text is pulled from the streamer
     # in the main thread. Adds timeout to the streamer to handle exceptions in the generation thread.