### Loading the Model

import torch
from transformers import AutoModelForVision2Seq, AutoProcessor, BitsAndBytesConfig

merged_model_save_path= "vector-institute/Llama3.2-Multimodal-Newsmedia-Bias-Detector"
# Load quantization config if used during saving
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

# Load the merged model
model = AutoModelForVision2Seq.from_pretrained(
    merged_model_save_path,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    quantization_config=bnb_config  
)
model.eval()

# Load the processor
processor = AutoProcessor.from_pretrained(merged_model_save_path)

import pandas as pd
import os
from PIL import Image

# Path to your dataset CSV and image folder
# get our sample from here https://huggingface.co/vector-institute/Llama3.2-Multimodal-Newsmedia-Bias-Detector/tree/main/sampled-data
dataset_csv_path = 'sample_dataset.csv'
image_folder_path = 'sampled_images'

# Load the DataFrame
df = pd.read_csv(dataset_csv_path)

# Function to prepare samples from the DataFrame
def prepare_samples_from_dataframe(df, image_folder_path):
    samples = []
    for index, row in df.iterrows():
        unique_id = row['unique_id']  # Replace with the column that contains the image identifier
        text_content = row['first_paragraph']  # Replace with your text column name

        # Find the image file
        possible_extensions = ['jpg', 'jpeg', 'png']
        image_path = None
        for ext in possible_extensions:
            img_path = os.path.join(image_folder_path, f"{unique_id}.{ext}")
            if os.path.exists(img_path):
                image_path = img_path
                break

        if image_path is None:
            print(f"No image found for ID {unique_id}")
            continue

        image = Image.open(image_path).convert("RGB")
        max_size = (224, 224)
        image.thumbnail(max_size, Image.Resampling.LANCZOS)

        # Prepare the prompt text
        sample_text = (
            "Assess the text and image below for potential disinformation (try finding deliberately misleading or biased information) by identifying the presence of rhetorical techniques listed.\n"
            "If you find any of the listed rhetorical techniques, then the article is likely disinformation; if not, it is likely not disinformation.\n\n"
            "Rhetorical Techniques Checklist:\n"
            "- Emotional Appeal: Uses language that intentionally invokes extreme emotions like fear or anger, aiming to distract from lack of factual backing.\n"
            "- Exaggeration and Hyperbole: Makes claims that are unsupported by evidence, or presents normal situations as extraordinary to manipulate perceptions.\n"
            "- Bias and Subjectivity: Presents information in a way that unreasonably favors one perspective, omitting key facts that might provide balance.\n"
            "- Repetition: Uses repeated messaging of specific points or misleading statements to embed a biased viewpoint in the reader's mind.\n"
            "- Specific Word Choices: Employs emotionally charged or misleading terms to sway opinions subtly, often in a manipulative manner.\n"
            "- Appeals to Authority: References authorities who lack relevant expertise or cites sources that do not have the credentials to be considered authoritative in the context.\n"
            "- Lack of Verifiable Sources: Relies on sources that either cannot be verified or do not exist, suggesting a fabrication of information.\n"
            "- Logical Fallacies: Engages in flawed reasoning such as circular reasoning, strawman arguments, or ad hominem attacks that undermine logical debate.\n"
            "- Conspiracy Theories: Propagates theories that lack proof and often contain elements of paranoia or implausible scenarios as facts.\n"
            "- Inconsistencies and Factual Errors: Contains multiple contradictions or factual inaccuracies that are easily disprovable, indicating a lack of concern for truth.\n"
            "- Selective Omission: Deliberately leaves out crucial information that is essential for a fair understanding of the topic, skewing perception.\n"
            "- Manipulative Framing: Frames issues in a way that leaves out alternative perspectives or possible explanations, focusing only on aspects that support a biased narrative.\n\n"
            f"{text_content}\n\n"
            "Please **only** provide your answer in the format: 'Classification: Likely' or 'Classification: Unlikely'. Do not include any additional text or explanation."
        )

        sample = {
            "unique_id": unique_id,
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": sample_text,
                        },
                        {
                            "type": "image",
                            "image": image,
                        }
                    ],
                }
            ]
        }

        samples.append(sample)

    return samples

import re

def extract_assistant_reply(generated_text):
    lines = generated_text.strip().split('\n')
    for idx, line in enumerate(lines):
        if line.strip().lower() == 'assistant':
            assistant_reply = '\n'.join(lines[idx+1:]).strip()
            return assistant_reply
    return lines[-1].strip()

def extract_classification(assistant_reply):
    match = re.search(r'Classification:\s*(Likely|Unlikely)', assistant_reply, re.IGNORECASE)
    if match:
        return match.group(1).capitalize()
    else:
        # Attempt to find 'Likely' or 'Unlikely' anywhere in the reply
        match = re.search(r'\b(Likely|Unlikely)\b', assistant_reply, re.IGNORECASE)
        if match:
            return match.group(1).capitalize()
        else:
            return 'Unknown'

def generate_prediction(sample):
    # Remove the model.to(device) line
    # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    # model.to(device)

    # Prepare the input texts and images
    texts = processor.apply_chat_template(sample["messages"], tokenize=False)
    image_input = sample["messages"][0]['content'][1]['image']

    # Prepare the inputs without moving them to a device
    inputs = processor(text=texts, images=image_input, return_tensors="pt", padding=True)

    # Ensure that the inputs are on the same device as the model's embeddings
    # This is handled internally by the model when using device_map="auto"

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=20,
            do_sample=False,
            num_beams=1,
        )

    generated_texts = processor.batch_decode(outputs, skip_special_tokens=True)
    assistant_reply = extract_assistant_reply(generated_texts[0])

    return assistant_reply

# Prepare the samples from the DataFrame
samples = prepare_samples_from_dataframe(df, image_folder_path)

# Limit to a subset if desired
samples_to_infer = samples[:900]  # For example, take the first 5 samples

# Run inference and collect results
results = []

for sample in samples_to_infer:
    assistant_reply = generate_prediction(sample)
    predicted_label = extract_classification(assistant_reply)
        # Get the first_paragraph from the original DataFrame
    original_row = df[df['unique_id'] == sample['unique_id']].iloc[0]
    first_paragraph = original_row['first_paragraph']
    # Collect results
    result = {
    'unique_id': sample['unique_id'],
    'first_paragraph': first_paragraph,  # Add this line
    'assistant_reply': assistant_reply,
    'predicted_label': predicted_label,
    }
    results.append(result)

    # Display the results
    print(f"Sample ID: {sample['unique_id']}")
    print(f"text: {first_paragraph}")
    print("Assistant's Reply:")
    print(assistant_reply)
    #print(f"Predicted Label: {predicted_label}")
    print("-" * 50)

results_df = pd.DataFrame(results)
# Save to CSV if desired
results_df.to_csv('llama-vision-ift-inference_results.csv', index=False)

Assistant's Reply:
Classification: Likely