| 118 | self.model=LlamaForCausalLM.from_pretrained(self.path,device_map="auto") |
| 119 | |
| 120 | def api(self,prompt): |
| 121 | inputs = self.tokenizer(prompt, return_tensors="pt", padding=True).to("cuda") |
| 122 | # Generate |
| 123 | outputs = self.model.generate(inputs.input_ids, max_length=2048, use_cache=True) |
| 124 | # return self.tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] |
| 125 | return self.tokenizer.batch_decode(outputs, skip_special_tokens=True) |
| 126 | |
| 127 | def parse_args(): |
| 128 | def list_of_strings(arg): |