Write a python script to do the following:
  - Use daft to read df from `s3://smoosense-demo/datasets/Rapidata/text-image-pair.parquet`
  - Take first 10 rows.
  - Download `image_url` and save to another column `image_bytes`
  - Define an UDF for misalignment:
      - Give `image_bytes` and `prompt` to a VLM, and expect misalignment score for each word. E.g. `[["A", 1], ["majestic", 0], ["horse", 5], ...]`
      - The higher score means the worse alignment. Range 0 to 5. 0 means perfect alignment. 5 means the opposite meaning.
      - Also give an overall misalignment_score for the prompt sentence. 
  - Save VLM response to columns `vlm_word_scores` and `vlm_misalignment_score`
  - Write output to parquet file. Do not include image_bytes.

import daft
import io
from PIL import Image
import base64
import json
import os
import requests
from anthropic import Anthropic

# Initialize Anthropic client
client = Anthropic(api_key=os.environ.get("ANTHROPIC_API_KEY"))

def download_image(url):
    """Download image from URL and return as bytes"""
    response = requests.get(url)
    response.raise_for_status()
    return response.content

def vlm_score_misalignment(image_bytes, prompt):
    """
    Use Claude Vision to analyze image-text misalignment.
    Returns a tuple: (word_scores_json, overall_score)
    """
    try:
        # Convert image bytes to base64
        image_b64 = base64.b64encode(image_bytes).decode('utf-8')

        # Determine image type
        image_type = "image/jpeg"
        try:
            img = Image.open(io.BytesIO(image_bytes))
            format_map = {
                'JPEG': 'image/jpeg',
                'PNG': 'image/png',
                'GIF': 'image/gif',
                'WEBP': 'image/webp'
            }
            image_type = format_map.get(img.format, 'image/jpeg')
        except:
            pass

        vlm_prompt = f"""Analyze the misalignment between this image and the following prompt: "{prompt}"

For each word in the prompt, assign a misalignment score from 0 to 5, where:
- 0 = perfect alignment with the image
- 1 = minor misalignment
- 2 = noticeable misalignment
- 3 = significant misalignment
- 4 = major misalignment
- 5 = opposite meaning or completely contradicts the image

Also provide an overall misalignment score (0-5) for the entire prompt sentence.

Return ONLY a valid JSON object in this exact format (no other text):
{{
  "word_scores": [["word1", score1], ["word2", score2], ...],
  "overall_score": score
}}

Example:
{{
  "word_scores": [["A", 1], ["majestic", 0], ["horse", 5], ["flying", 4]],
  "overall_score": 3.5
}}

Prompt to analyze: {prompt}"""

        message = client.messages.create(
            model="claude-sonnet-4-20250514",
            max_tokens=2048,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image",
                            "source": {
                                "type": "base64",
                                "media_type": image_type,
                                "data": image_b64,
                            },
                        },
                        {
                            "type": "text",
                            "text": vlm_prompt
                        }
                    ],
                }
            ],
        )

        # Extract response
        response_text = message.content[0].text.strip()

        # Try to extract JSON if there's extra text
        if not response_text.startswith('{'):
            # Try to find JSON object in the response
            start_idx = response_text.find('{')
            end_idx = response_text.rfind('}') + 1
            if start_idx != -1 and end_idx > start_idx:
                response_text = response_text[start_idx:end_idx]

        # Parse JSON
        result = json.loads(response_text)

        # Extract word scores and overall score
        word_scores_str = json.dumps(result.get("word_scores", []))
        overall_score = result.get("overall_score", 0)

        return word_scores_str, overall_score

    except Exception as e:
        print(f"Error in VLM scoring: {e}")
        return "[]", 0

def main():
    # Read parquet from S3 using daft
    print("Reading data from S3...")
    df = daft.read_parquet("s3://smoosense-demo/datasets/Rapidata/text-image-pair.parquet")

    df = df.limit(10)

    # Convert to pandas for easier processing
    df = df.to_pandas()

    print(f"Processing {len(df)} rows")
    print(f"Columns: {df.columns.tolist()}")

    # Download images
    print("Downloading images...")
    df['image_bytes'] = df['image_url'].apply(download_image)

    # Apply VLM scoring
    print("Running VLM misalignment scoring...")
    results = df.apply(
        lambda row: vlm_score_misalignment(row['image_bytes'], row['prompt']),
        axis=1
    )

    # Unpack results into separate columns
    df['vlm_word_scores'] = results.apply(lambda x: x[0])
    df['vlm_misalignment_score'] = results.apply(lambda x: x[1])

    # Drop image_bytes and keep only required columns
    output_df = df.drop(columns=['image_bytes'])

    # Write to parquet
    output_path = "output_misalignment_scores.parquet"
    print(f"Writing output to {output_path}...")
    output_df.to_parquet(output_path, index=False)

    print("Done!")
    print(f"Output saved to: {output_path}")
    print(f"\nSample output:")
    print(output_df[['image_url', 'vlm_word_scores', 'vlm_misalignment_score']].head(3))

if __name__ == "__main__":
    main()