| 111 | return processed_prompt |
| 112 | |
| 113 | class API: |
| 114 | def __init__(self, model_path): |
| 115 | self.path = model_path |
| 116 | self.tokenizer=AutoTokenizer.from_pretrained(self.path) |
| 117 | self.tokenizer.pad_token=self.tokenizer.eos_token |
| 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): |
nothing calls this directly
no outgoing calls
no test coverage detected