is there any issues with this code 
"
import os
import pandas as pd
import numpy as np
from PIL import Image
import requests
from io import BytesIO
import torch
import torchvision.transforms as transforms
from torch.utils.data import Dataset, DataLoader
from tqdm import tqdm
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error, f1_score, precision_score, recall_score
import joblib
import timm
import logging
from transformers import ViTImageProcessor, ViTModel
from catboost import CatBoostRegressor, Pool

# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# Use GPU if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logging.info(f"Using device: {device}")

def download_image(image_link, save_path):
    try:
        response = requests.get(image_link, timeout=10)
        img = Image.open(BytesIO(response.content)).convert('RGB')
        img.save(save_path)
        return img
    except Exception as e:
        logging.warning(f"Error processing image: {image_link} - {e}")
        img = Image.new('RGB', (224, 224), color='white')
        img.save(save_path)
        return img

class ProductImageDataset(Dataset):
    def __init__(self, image_paths, labels=None, transform=None):
        self.image_paths = image_paths
        self.labels = labels
        self.transform = transform

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        img = Image.open(self.image_paths[idx]).convert('RGB')
        if self.transform:
            img = self.transform(img)
        
        if self.labels is not None:
            return img, self.labels[idx]
        return img

def extract_features(image_paths, batch_size=32):
    image_processor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224-in21k')
    model = ViTModel.from_pretrained('google/vit-base-patch16-224-in21k')
    model = model.to(device)
    model.eval()

    dataset = ProductImageDataset(image_paths, transform=lambda x: image_processor(x, return_tensors="pt")["pixel_values"][0])
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=4, pin_memory=True)

    features = []
    with torch.no_grad():
        for batch in tqdm(dataloader, desc="Extracting features"):
            batch = batch.to(device)
            outputs = model(pixel_values=batch)
            features.append(outputs.last_hidden_state[:, 0, :].cpu().numpy())

    return np.vstack(features)

def train_model(X, y, entity_names, test_features, test_entity_names):
    # Remove NaN values
    mask = ~np.isnan(y)
    X = X[mask]
    y = y[mask]
    entity_names = entity_names[mask]

    X_train, X_val, y_train, y_val, entity_train, entity_val = train_test_split(X, y, entity_names, test_size=0.2, random_state=42)

    # Create CatBoost Pool objects
    train_pool = Pool(X_train, y_train, cat_features=[X_train.shape[1]])
    val_pool = Pool(X_val, y_val, cat_features=[X_val.shape[1]])

    logging.info("Training CatBoost model...")
    model = CatBoostRegressor(
        iterations=2000,
        learning_rate=0.05,
        depth=8,
        l2_leaf_reg=3,
        loss_function='RMSE',
        verbose=100,
        early_stopping_rounds=50
    )
    
    try:
        model.fit(train_pool, eval_set=val_pool, plot=True)
    except Exception as e:
        logging.error(f"Error during model training: {e}")
        return None, None

    y_pred = model.predict(X_val)
    mse = mean_squared_error(y_val, y_pred)
    logging.info(f"Validation MSE: {mse}")

    # Calculate F1 score, precision, and recall
    threshold = 0.5  # You may need to adjust this threshold
    y_pred_binary = (y_pred > threshold).astype(int)
    y_val_binary = (y_val > threshold).astype(int)
    f1 = f1_score(y_val_binary, y_pred_binary)
    precision = precision_score(y_val_binary, y_pred_binary)
    recall = recall_score(y_val_binary, y_pred_binary)
    
    logging.info(f"Validation F1 Score: {f1}")
    logging.info(f"Validation Precision: {precision}")
    logging.info(f"Validation Recall: {recall}")

    test_features_with_entity = np.column_stack((test_features, pd.get_dummies(test_entity_names)))
    test_predictions = model.predict(test_features_with_entity)
    
    return model, test_predictions

def format_prediction(entity_name, value):
    units = {
        'width': 'centimetre', 'height': 'centimetre', 'depth': 'centimetre',
        'item_weight': 'gram', 'voltage': 'volt', 'wattage': 'watt',
        'maximum_weight_recommendation': 'kilogram', 'item_volume': 'litre'
    }
    if value <= 0:
        return ""
    return f"{value:.2f} {units.get(entity_name, '')}" if entity_name in units else ""

def main():
    try:
        # Load and preprocess data
        train = pd.read_csv('train.csv')
        test = pd.read_csv('test.csv')

        train['entity_value_numeric'] = pd.to_numeric(train['entity_value'].str.split().str[0], errors='coerce')

        # Create directories for images
        os.makedirs('train_images', exist_ok=True)
        os.makedirs('test_images', exist_ok=True)

        # Download and save images
        logging.info("Downloading training images...")
        train['image_path'] = train['image_link'].apply(lambda x: os.path.join('train_images', x.split('/')[-1]))
        for idx, row in tqdm(train.iterrows(), total=len(train)):
            download_image(row['image_link'], row['image_path'])

        logging.info("Downloading test images...")
        test['image_path'] = test['image_link'].apply(lambda x: os.path.join('test_images', x.split('/')[-1]))
        for idx, row in tqdm(test.iterrows(), total=len(test)):
            download_image(row['image_link'], row['image_path'])

        # Extract features
        logging.info("Extracting features for training data...")
        train_features = extract_features(train['image_path'].tolist())
        logging.info("Extracting features for test data...")
        test_features = extract_features(test['image_path'].tolist())

        # Add entity name as a categorical feature
        train_features_with_entity = np.column_stack((train_features, pd.get_dummies(train['entity_name'])))
        
        # Train model and get predictions
        model, test_predictions = train_model(train_features_with_entity, train['entity_value_numeric'].values, 
                                              train['entity_name'].values, test_features, test['entity_name'].values)

        if model is None or test_predictions is None:
            logging.error("Model training failed. Exiting.")
            return

        # Format predictions
        test['prediction'] = test.apply(lambda row: format_prediction(row['entity_name'], test_predictions[row.name]), axis=1)

        # Save the predictions
        test[['index', 'prediction']].to_csv('test_out.csv', index=False)
        logging.info("Predictions saved to test_out.csv")

        # Save the model
        joblib.dump(model, 'catboost_model.joblib')
        logging.info("Model saved as catboost_model.joblib")

    except Exception as e:
        logging.error(f"An error occurred: {e}")

if __name__ == "__main__":
    main()
"
here is my problem statement
"
# ML Challenge Problem Statement

## Feature Extraction from Images

In this hackathon, the goal is to create a machine learning model that extracts entity values from images. This capability is crucial in fields like healthcare, e-commerce, and content moderation, where precise product information is vital. As digital marketplaces expand, many products lack detailed textual descriptions, making it essential to obtain key details directly from images. These images provide important information such as weight, volume, voltage, wattage, dimensions, and many more, which are critical for digital stores.

### Data Description: 

The dataset consists of the following columns: 

1. **index:** An unique identifier (ID) for the data sample
2. **image_link**: Public URL where the product image is available for download. Example link - https://m.media-amazon.com/images/I/71XfHPR36-L.jpg
    To download images use `download_images` function from `src/utils.py`. See sample code in `src/test.ipynb`.
3. **group_id**: Category code of the product
4. **entity_name:** Product entity name. For eg: “item_weight” 
5. **entity_value:** Product entity value. For eg: “34 gram” 
    Note: For test.csv, you will not see the column `entity_value` as it is the target variable.

### Output Format:

The output file should be a csv with 2 columns:

1. **index:** The unique identifier (ID) of the data sample. Note the index should match the test record index.
2. **prediction:** A string which should have the following format: “x unit” where x is a float number in standard formatting and unit is one of the allowed units (allowed units are mentioned in the Appendix). The two values should be concatenated and have a space between them. For eg: “2 gram”, “12.5 centimetre”, “2.56 ounce” are valid. Few invalid cases: “2 gms”, “60 ounce/1.7 kilogram”, “2.2e2 kilogram” etc.
    Note: Make sure to output a prediction for all indices. If no value is found in the image for any test sample, return empty string, i.e, `“”`. If you have less/more number of output samples in the output file as compared to test.csv, your output won’t be evaluated. 

### File Descriptions:

*source files*

1. **src/sanity.py**: Sanity checker to ensure that the final output file passes all formatting checks. Note: the script will not check if less/more number of predictions are present compared to the test file. See sample code in `src/test.ipynb` 
2. **src/utils.py**: Contains helper functions for downloading images from the image_link.
3. **src/constants.py:** Contains the allowed units for each entity type.
4. **sample_code.py:** We also provided a sample dummy code that can generate an output file in the given format. Usage of this file is optional. 

*Dataset files*

1. **dataset/train.csv**: Training file with labels (`entity_value`).
2. **dataset/test.csv**: Test file without output labels (`entity_value`). Generate predictions using your model/solution on this file's data and format the output file to match sample_test_out.csv (Refer the above section "Output Format")
3. **dataset/sample_test.csv**: Sample test input file.
4. **dataset/sample_test_out.csv**: Sample outputs for sample_test.csv. The output for test.csv must be formatted in the exact same way. Note: The predictions in the file might not be correct

### Constraints

1. You will be provided with a sample output file and a sanity checker file. Format your output to match the sample output file exactly and pass it through the sanity checker to ensure its validity. Note: If the file does not pass through the sanity checker, it will not be evaluated. You should recieve a message like `Parsing successfull for file: ...csv` if the output file is correctly formatted.

2. You are given the list of allowed units in constants.py and also in Appendix. Your outputs must be in these units. Predictions using any other units will be considered invalid during validation.

### Evaluation Criteria

Submissions will be evaluated based on F1 score, which are standard measures of prediction accuracy for classification and extraction problems.

Let GT = Ground truth value for a sample and OUT be output prediction from the model for a sample. Then we classify the predictions into one of the 4 classes with the following logic: 

1. *True Positives* - If OUT != `""` and GT != `""` and OUT == GT
2. *False Positives* - If OUT != `""` and GT != `""` and OUT != GT
3. *False Positives* - If OUT != `""` and GT == `""`
4. *False Negatives* - If OUT == `""` and GT != `""`
5. *True Negatives* - If OUT == `""` and GT == `""` 

Then, F1 score = 2*Precision*Recall/(Precision + Recall) where:

1. Precision = True Positives/(True Positives + False Positives)
2. Recall = True Positives/(True Positives + False Negatives)

### Submission File

Upload a test_out.csv file in the Portal with the exact same formatting as sample_test_out.csv

### Appendix

```
entity_unit_map = {
  "width": {
    "centimetre",
    "foot",
    "millimetre",
    "metre",
    "inch",
    "yard"
  },
  "depth": {
    "centimetre",
    "foot",
    "millimetre",
    "metre",
    "inch",
    "yard"
  },
  "height": {
    "centimetre",
    "foot",
    "millimetre",
    "metre",
    "inch",
    "yard"
  },
  "item_weight": {
    "milligram",
    "kilogram",
    "microgram",
    "gram",
    "ounce",
    "ton",
    "pound"
  },
  "maximum_weight_recommendation": {
    "milligram",
    "kilogram",
    "microgram",
    "gram",
    "ounce",
    "ton",
    "pound"
  },
  "voltage": {
    "millivolt",
    "kilovolt",
    "volt"
  },
  "wattage": {
    "kilowatt",
    "watt"
  },
  "item_volume": {
    "cubic foot",
    "microlitre",
    "cup",
    "fluid ounce",
    "centilitre",
    "imperial gallon",
    "pint",
    "decilitre",
    "litre",
    "millilitre",
    "quart",
    "cubic inch",
    "gallon"
  }
}
```

"

and based on it give me the complete code