Move input_ids to CUDA about jsonformer

[OmarMohammed88]

can you share the code u used?

[anujnayyar1]

from jsonformer import Jsonformer
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

tokenizer = AutoTokenizer.from_pretrained("StabilityAI/stablelm-base-alpha-7b")
model = AutoModelForCausalLM.from_pretrained("StabilityAI/stablelm-base-alpha-7b")
model.half().cuda()

json_schema = {
    "type": "object",
    "properties": {
        "name": {"type": "string"},
        "age": {"type": "number"},
        "is_student": {"type": "boolean"},
        "courses": {
            "type": "array",
            "items": {"type": "string"}
        }
    }
}

prompt = "Generate a person's information based on the following schema:"
jsonformer = Jsonformer(model, tokenizer, json_schema, prompt)
generated_data = jsonformer()

print(generated_data)

To alleviate this problem I have tried loading through:
model = AutoModelForCausalLM.from_pretrained("StabilityAI/stablelm-base-alpha-7b", device_map="auto", torch_dtype=torch.float16)

Which seems to load better but then I run into the following:

IndexError                                Traceback (most recent call last)
Cell In[10], line 23
     21 prompt = "Generate a person's information based on the following schema:"
     22 jsonformer = Jsonformer(model, tokenizer, json_schema, prompt)
---> 23 generated_data = jsonformer()
     25 print(generated_data)

File /opt/conda/lib/python3.10/site-packages/jsonformer/main.py:188, in Jsonformer.__call__(self)
    185 def __call__(self) -> Dict[str, Any]:
    186     self.value = {}
--> 188     generated_data = self.generate_object(
    189         self.json_schema["properties"], self.value
    190     )
    191     return generated_data

File /opt/conda/lib/python3.10/site-packages/jsonformer/main.py:114, in Jsonformer.generate_object(self, properties, obj)
    109 def generate_object(
    110     self, properties: Dict[str, Any], obj: Dict[str, Any]
    111 ) -> Dict[str, Any]:
    112     # self.debug("[generate_object] properties", properties)
    113     for key, schema in properties.items():
--> 114         obj[key] = self.generate_value(schema, obj, key)
    115     return obj

File /opt/conda/lib/python3.10/site-packages/jsonformer/main.py:136, in Jsonformer.generate_value(self, schema, obj, key)
    134     new_array = []
    135     obj[key] = new_array
--> 136     return self.generate_array(schema["items"], new_array)
    137 elif schema_type == "object":
    138     new_obj = {}

File /opt/conda/lib/python3.10/site-packages/jsonformer/main.py:146, in Jsonformer.generate_array(self, item_schema, obj)
    144 def generate_array(self, item_schema: Dict[str, Any], obj: Dict[str, Any]) -> list:
    145     for _ in range(self.max_array_length):
--> 146         element = self.generate_value(item_schema, obj)
    147         obj[-1] = element
    149         obj.append(self.generation_marker)

File /opt/conda/lib/python3.10/site-packages/jsonformer/main.py:131, in Jsonformer.generate_value(self, schema, obj, key)
    129     return self.generate_boolean()
    130 elif schema_type == "string":
--> 131     obj[key if key else -1] = self.generation_marker
    132     return self.generate_string()
    133 elif schema_type == "array":

IndexError: list assignment index out of range

[1rgs]

@anujnayyar1 this should be fixed in the latest release, thanks!