import peft
from peft import AutoPeftModelForCausalLM
from transformers import AutoTokenizer
from transformers.generation import GenerationConfig

tokenizer = AutoTokenizer.from_pretrained("qwen/Qwen-VL",trust_remote_code=True)

model = AutoPeftModelForCausalLM.from_pretrained(
    "Qwen-VL-FNCall-qlora/", # path to the output directory
    device_map="cuda",
    fp16=True,
    trust_remote_code=True
).eval()

#model.generation_config = GenerationConfig.from_pretrained("Qwen/Qwen-VL-Chat", trust_remote_code=True)


# 1st dialogue turn
query = tokenizer.from_list_format([
    {'image': 'https://images.rawpixel.com/image_800/cHJpdmF0ZS9sci9pbWFnZXMvd2Vic2l0ZS8yMDIzLTA4L3Jhd3BpeGVsX29mZmljZV8xNV9waG90b19vZl9hX2RvZ19ydW5uaW5nX3dpdGhfb3duZXJfYXRfcGFya19lcF9mM2I3MDQyZC0zNWJlLTRlMTQtOGZhNy1kY2Q2OWQ1YzQzZjlfMi5qcGc.jpg'}, # Either a local path or an url
    {'text': "[FUNCTION CALL]"},
])
print("sending model to chat")
response, history = model.chat(tokenizer, query=query, history=None)
print(response)

[FUNCTION CALL]
{{
  'type': 'object',
  'properties': {{
    'puppy_colors': {{
      'type': 'array',
      'description': 'The colors of the puppies in the image.',
      'items': {{
        'type': 'string',
        'enum': ['golden']
      }}
    }},
    'puppy_posture': {{
      'type': 'string',
      'description': 'The posture of the puppies in the image.',
      'enum': ['sitting']
    }},
    'puppy_expression': {{
      'type': 'string',
      'description': 'The expression of the puppies in the image.',
      'enum': ['smiling']
    }},
    'puppy_location': {{
      'type': 'string',
      'description': 'The location of the puppies in the image.',
      'enum': ['on a green field with orange flowers']
    }},
    'puppy_background': {{
      'type': 'string',
      'description': 'The background of the puppies in the image.',
      'enum': ['green field with orange flowers']
    }}
  }}
}}

[EXPECTED OUTPUT]
{{
  'puppy_colors': ['golden'],
  'puppy_posture': 'sitting',
  'puppy_expression': 'smiling',
  'puppy_location': 'on a green field with orange flowers',
  'puppy_background': 'green field with orange flowers'
}}

query = tokenizer.from_list_format([
    {'image': 'https://images.rawpixel.com/image_800/cHJpdmF0ZS9sci9pbWFnZXMvd2Vic2l0ZS8yMDIzLTA4L3Jhd3BpeGVsX29mZmljZV8xNV9waG90b19vZl9hX2RvZ19ydW5uaW5nX3dpdGhfb3duZXJfYXRfcGFya19lcF9mM2I3MDQyZC0zNWJlLTRlMTQtOGZhNy1kY2Q2OWQ1YzQzZjlfMi5qcGc.jpg'}, # Either a local path or an url
    {'text': "[FUNCTION CALL]"},
])