Fix some issues

Author: shubhamugare

requirements.txt

@@ -1,3 +1,4 @@
+accelerate
 fire
 interegular
 regex==2023.8.8

syncode/language_model.py

@@ -176,7 +176,11 @@ def get_tokenized_input(self, prompt: Union[str, list], batch_size: int):
             raise ValueError("Prompt should be either a string or a list! It is currently of type: "+str(type(prompt)))
 
         input_batch = [prompt_str for _ in range(batch_size)]
-        inputs = self.tokenizer(input_batch, return_tensors="pt").to(self.model.device)
+        inputs = self.tokenizer(
+            input_batch, 
+            return_tensors="pt",
+            pad_to_multiple_of=8,
+            ).to(self.model.device)
 
         return inputs
 

tests/test_language_model.py

@@ -56,19 +56,17 @@ def get_vocab(self) -> Dict[str, int]:
         return {v: i for i, v in enumerate(self.vocab)}
 
 class TestHuggingFaceModel(unittest.TestCase):
-    @unittest.skip("Only for local testing")
     def test_stop_word(self):
         torch.manual_seed(0)
-        syncode = Syncode(model="microsoft/phi-2", mode='original')
+        syncode = Syncode(model="microsoft/phi-2", mode='original', device='cpu')
         prompt = "Generate a json for the country nigeria.\n```json\n"
         stop_words = ["```"]
         output = syncode.infer(prompt, stop_words=stop_words)[0]
         assert output.endswith('```')
 
-    @unittest.skip("Only for local testing")
     def test_stop_word2(self):
         torch.manual_seed(0)
-        syncode = Syncode(model="microsoft/phi-2", mode='original')
+        syncode = Syncode(model="microsoft/phi-2", mode='original', device='cpu')
         prompt = "def add(a, b):\n"
         stop_words = ["\n\n"]
         output = syncode.infer(prompt, stop_words=stop_words)[0]