Spaces:
Runtime error
Runtime error
| from transformers import AutoTokenizer, AutoModelForCausalLM | |
| class Petals: | |
| """Petals Bloom models.""" | |
| def __init__( | |
| self, | |
| model_name="bigscience/bloom-petals", | |
| temperature=0.7, | |
| max_new_tokens=256, | |
| top_p=0.9, | |
| top_k=None, | |
| do_sample=True, | |
| max_length=None | |
| ): | |
| self.model_name = model_name | |
| self.temperature = temperature | |
| self.max_new_tokens = max_new_tokens | |
| self.top_p = top_p | |
| self.top_k = top_k | |
| self.do_sample = do_sample | |
| self.max_length = max_length | |
| self.tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| self.model = AutoModelForCausalLM.from_pretrained(model_name) | |
| def _default_params(self): | |
| """Get the default parameters for calling Petals API.""" | |
| return { | |
| "temperature": self.temperature, | |
| "max_new_tokens": self.max_new_tokens, | |
| "top_p": self.top_p, | |
| "top_k": self.top_k, | |
| "do_sample": self.do_sample, | |
| "max_length": self.max_length, | |
| } | |
| def generate(self, prompt): | |
| """Generate text using the Petals API.""" | |
| params = self._default_params() | |
| inputs = self.tokenizer(prompt, return_tensors="pt")["input_ids"] | |
| outputs = self.model.generate(inputs, **params) | |
| return self.tokenizer.decode(outputs[0]) |