Upload prompt.py

competition/prompt.py

@@ -0,0 +1,161 @@
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments
+from datasets import Dataset
+import pandas as pd
+from sklearn.model_selection import train_test_split
+from peft import get_peft_model, LoraConfig, TaskType
+import evaluate
+import numpy as np
+from tqdm import tqdm
+
+# Load the dataset
+file_path = 'train_en.csv'
+dataset = pd.read_csv(file_path)
+
+# Map labels to expected responses
+label_mapping = {
+    "Yes": 0,
+    "No": 1,
+    "It doesn't matter": 2,
+    "Unimportant": 2,
+    "Incorrect questioning": 3,
+    "Correct answers": 4
+}
+
+# Apply label mapping
+dataset['label'] = dataset['label'].map(label_mapping)
+
+# Handle NaN values: Drop rows where label is NaN
+dataset = dataset.dropna(subset=['label'])
+
+# Ensure labels are integers
+dataset['label'] = dataset['label'].astype(int)
+
+# Format puzzle, truth, text into the prompt
+prompt_template = """You are the host of a situational guessing game. The rules of the game are as follows:
+
+1. Participants will receive a riddle that describes a simple yet difficult to understand event.
+2. The host knows the answer, which is the solution to the riddle.
+3. Participants can ask any closed-ended questions to uncover the truth of the event.
+4. For each question, the host will respond with one of the following five options based on the actual situation: Yes, No, Unimportant, Correct answer, or Incorrect questioning. The criteria for each response are as follows:
+   - If the riddle and answer can provide an answer to the question, respond with: Yes or No
+   - If the riddle and answer cannot directly or indirectly infer an answer to the question, respond with: Unimportant
+   - If the participant's question is not a closed-ended question or is difficult to understand, respond with: Incorrect questioning
+   - If the participant's question essentially reveals the truth of the answer, respond with: Correct answer
+5. The response must not include any additional information, nor should any word be omitted from the options. For example, "No" cannot be abbreviated to "N".
+
+Please strictly follow these rules when answering the participant's questions.
+
+Riddle: {}
+Answer: {}
+Participant's question: {}
+"""
+
+dataset['combined_text'] = dataset.apply(
+    lambda row: prompt_template.format(row['puzzle'], row['truth'], row['text']),
+    axis=1
+)
+
+# Split the dataset into training and validation sets
+train_df, val_df = train_test_split(dataset, test_size=0.2, random_state=42)
+
+# Convert the dataframes to datasets
+train_dataset = Dataset.from_pandas(train_df)
+val_dataset = Dataset.from_pandas(val_df)
+
+# Load the tokenizer and model
+model_name = "meta-llama/Meta-Llama-3-8B"  # Replace with the actual model name
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=5)
+
+# Add a padding token if it's not already present
+if tokenizer.pad_token is None:
+    tokenizer.add_special_tokens({'pad_token': tokenizer.eos_token})
+    model.resize_token_embeddings(len(tokenizer))
+    tokenizer.pad_token = tokenizer.eos_token  # Set the padding token explicitly
+
+# Ensure the padding token is set correctly in the model configuration
+model.config.pad_token_id = tokenizer.pad_token_id
+
+# Tokenize the data
+def tokenize_function(examples):
+    return tokenizer(examples['combined_text'], truncation=True, padding='max_length', max_length=512)
+
+train_dataset = train_dataset.map(tokenize_function, batched=True, num_proc=4)  # Use multiprocessing
+val_dataset = val_dataset.map(tokenize_function, batched=True, num_proc=4)
+
+# Set the format for PyTorch
+train_dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'label'])
+val_dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'label'])
+
+# Define LoRA configuration
+lora_config = LoraConfig(
+    task_type=TaskType.SEQ_CLS,
+    r=16,
+    lora_alpha=16,
+    target_modules=["q_proj", "v_proj"],
+    lora_dropout=0.05,
+    bias="none"
+)
+
+# Apply LoRA to the model
+model = get_peft_model(model, lora_config)
+model.print_trainable_parameters()
+
+# Training arguments
+training_args = TrainingArguments(
+    output_dir='./results',
+    learning_rate=1e-4,
+    lr_scheduler_type="linear",
+    warmup_ratio=0.1,
+    max_grad_norm=0.3,
+    per_device_train_batch_size=8,  # Increase batch size if memory allows
+    per_device_eval_batch_size=8,
+    num_train_epochs=3,
+    weight_decay=0.001,
+    evaluation_strategy="epoch",
+    save_strategy="epoch",
+    load_best_model_at_end=True,
+    report_to="wandb",
+    fp16=True,
+    gradient_checkpointing=True,
+    gradient_accumulation_steps=2,  # Adjust based on memory constraints
+    dataloader_num_workers=4,
+    logging_steps=100,
+    save_total_limit=2,
+)
+
+def compute_metrics(eval_pred):
+    precision_metric = evaluate.load("precision")
+    recall_metric = evaluate.load("recall")
+    f1_metric = evaluate.load("f1")
+    accuracy_metric = evaluate.load("accuracy")
+
+    logits, labels = eval_pred
+    predictions = np.argmax(logits, axis=-1)
+
+    precision = precision_metric.compute(predictions=predictions, references=labels, average="weighted")["precision"]
+    recall = recall_metric.compute(predictions=predictions, references=labels, average="weighted")["recall"]
+    f1 = f1_metric.compute(predictions=predictions, references=labels, average="weighted")["f1"]
+    accuracy = accuracy_metric.compute(predictions=predictions, references=labels)["accuracy"]
+    
+    return {"precision": precision, "recall": recall, "f1-score": f1, 'accuracy': accuracy}
+
+# Initialize the Trainer
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=train_dataset,
+    eval_dataset=val_dataset,
+    compute_metrics=compute_metrics
+)
+
+# Train the model with progress bar
+trainer.train()
+
+# Save the model
+model.save_pretrained('trained_llama_model')
+tokenizer.save_pretrained('trained_llama_model')
+
+# Evaluate the model with progress bar
+eval_results = trainer.evaluate()
+print(eval_results)