Update fine_tuning.py

fine_tuning.py

@@ -41,62 +41,50 @@ def prepare_data(input_texts, target_texts):
     targets = tokenizer(target_texts, max_length=512, truncation=True, padding="max_length")
     return {"input_ids": inputs["input_ids"], "attention_mask": inputs["attention_mask"], "labels": targets["input_ids"]}
 
-def paraphrase_with_model(text, model, tokenizer):
-    prompt = "Create a detailed, structured teaching transcript from the following text: " + text
-    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=512)
-    output_ids = model.generate(
-        inputs["input_ids"],
-        do_sample=True,
-        top_k=40,
-        top_p=0.9,
-        temperature=0.8,
-        max_length=300,
-        no_repeat_ngram_size=3,
-        early_stopping=True
-    )
-    return tokenizer.decode(output_ids[0], skip_special_tokens=True)
-
-
-model_name = "t5-base"
-tokenizer = T5Tokenizer.from_pretrained(model_name)
-model = T5ForConditionalGeneration.from_pretrained(model_name)
-
-try:
-    logger.info("Reading and cleaning prompts.")
-    input_texts, target_texts = read_prompts("prompts.txt")
-    input_texts_cleaned = [clean_text(text) for text in input_texts]
-    target_texts_cleaned = [clean_text(text) for text in target_texts]
-
-    logger.info("Splitting dataset into training and validation sets.")
-    train_texts, val_texts, train_labels, val_labels = train_test_split(input_texts_cleaned, target_texts_cleaned, test_size=0.1)
-
-    logger.info("Preparing datasets for training.")
-    train_dataset = Dataset.from_dict(prepare_data(train_texts, train_labels))
-    val_dataset = Dataset.from_dict(prepare_data(val_texts, val_labels))
-
-    training_args = TrainingArguments(
-        output_dir="./results",
-        evaluation_strategy="steps",
-        learning_rate=5e-5,
-        per_device_train_batch_size=4,
-        num_train_epochs=3,
-        save_steps=500,
-        logging_dir="./logs",
-        logging_steps=10
-    )
-
-    logger.info("Starting model training.")
-    trainer = Trainer(
-        model=model,
-        args=training_args,
-        train_dataset=train_dataset,
-        eval_dataset=val_dataset
-    )
-    trainer.train()
-
-    logger.info("Saving fine-tuned model.")
-    model.save_pretrained("./fine_tuned_model")
-    tokenizer.save_pretrained("./fine_tuned_model")
-
-except Exception as e:
-    logger.error(f"An error occurred during fine-tuning: {str(e)}")
+# Fine-tuning
+def fine_tune_model():
+    model_name = "t5-base"
+    tokenizer = T5Tokenizer.from_pretrained(model_name)
+    model = T5ForConditionalGeneration.from_pretrained(model_name)
+
+    try:
+        logger.info("Reading and cleaning prompts.")
+        input_texts, target_texts = read_prompts("prompts.txt")
+        input_texts_cleaned = [clean_text(text) for text in input_texts]
+        target_texts_cleaned = [clean_text(text) for text in target_texts]
+
+        logger.info("Splitting dataset into training and validation sets.")
+        train_texts, val_texts, train_labels, val_labels = train_test_split(input_texts_cleaned, target_texts_cleaned, test_size=0.1)
+
+        logger.info("Preparing datasets for training.")
+        train_dataset = Dataset.from_dict(prepare_data(train_texts, train_labels, tokenizer))
+        val_dataset = Dataset.from_dict(prepare_data(val_texts, val_labels, tokenizer))
+
+        training_args = TrainingArguments(
+            output_dir="./results",
+            evaluation_strategy="steps",
+            learning_rate=5e-5,
+            per_device_train_batch_size=4,
+            num_train_epochs=3,
+            save_steps=500,
+            logging_dir="./logs",
+            logging_steps=10
+        )
+
+        logger.info("Starting model training.")
+        trainer = Trainer(
+            model=model,
+            args=training_args,
+            train_dataset=train_dataset,
+            eval_dataset=val_dataset
+        )
+        trainer.train()
+
+        logger.info("Saving fine-tuned model.")
+        model.save_pretrained("./fine_tuned_model")
+        tokenizer.save_pretrained("./fine_tuned_model")
+
+    except Exception as e:
+        logger.error(f"An error occurred during fine-tuning: {str(e)}")
+
+fine_tune_model()