Upload 3 files

all_in_one.py

@@ -1,115 +1,117 @@
-import os
-import torch
-import json
-from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments, DataCollatorForLanguageModeling, TrainerCallback
-from datasets import Dataset
-import matplotlib.pyplot as plt
-
-# Set Hugging Face token (replace with your actual token)
-os.environ["HF_TOKEN"] = "hf_XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX"  # Replace with your HF_TOKEN
-
-# Download model and tokenizer
-model_name = "Salesforce/codegen-350M-multi"
-local_model_path = "./codegen_model"
-tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=local_model_path)
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float32, cache_dir=local_model_path)
-
-# Set padding token
-tokenizer.pad_token = tokenizer.eos_token
-
-# Move model to CPU
-device = torch.device("cpu")
-model.to(device)
-
-# Load custom dataset from JSONL
-dataset_path = "./custom_dataset.jsonl"
-data = []
-with open(dataset_path, 'r', encoding='utf-8') as f:
-    for line in f:
-        data.append(json.loads(line.strip()))
-dataset = Dataset.from_list(data)
-
-# Tokenize dataset
-def tokenize_function(examples):
-    inputs = [f"{prompt}\n{code}" for prompt, code in zip(examples["prompt"], examples["code"])]
-    return tokenizer(inputs, truncation=True, padding="max_length", max_length=128)
-
-tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=["prompt", "code"])
-
-# Data collator for language modeling
-data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
-
-# Define training arguments
-training_args = TrainingArguments(
-    output_dir="./finetuned_codegen",
-    overwrite_output_dir=True,
-    num_train_epochs=3,
-    per_device_train_batch_size=1,
-    gradient_accumulation_steps=4,
-    save_steps=500,
-    save_total_limit=2,
-    logging_steps=100,
-    learning_rate=5e-5,
-    fp16=False,
-    no_cuda=True,
-    dataloader_pin_memory=False,
-)
-
-# Custom callback to store training loss
-class LossCallback(TrainerCallback):
-    def __init__(self):
-        self.losses = []
-
-    def on_log(self, args, state, control, logs=None, **kwargs):
-        if logs and "loss" in logs:
-            self.losses.append(logs["loss"])
-
-loss_callback = LossCallback()
-
-# Initialize Trainer
-trainer = Trainer(
-    model=model,
-    args=training_args,
-    train_dataset=tokenized_dataset,
-    data_collator=data_collator,
-    callbacks=[loss_callback],
-)
-
-# Start fine-tuning
-print("Starting fine-tuning...")
-trainer.train()
-
-# Save fine-tuned model
-model.save_pretrained("./finetuned_codegen")
-tokenizer.save_pretrained("./finetuned_codegen")
-
-# Plot training loss
-plt.plot(loss_callback.losses, label="Training Loss")
-plt.xlabel("Steps")
-plt.ylabel("Loss")
-plt.title("Fine-Tuning Loss Curve")
-plt.legend()
-plt.savefig("./finetuned_codegen/loss_plot.png")
-plt.show()
-
-print("Fine-tuning completed. Model saved to ./finetuned_codegen. Loss plot saved to ./finetuned_codegen/loss_plot.png")
-
-# Test fine-tuned model
-print("\nTesting fine-tuned model...")
-prompts = [
-    "Write a Python program to print 'Hello, World!'"
-]
-
-for prompt in prompts:
-    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=128).to(device)
-    outputs = model.generate(
-        **inputs,
-        max_length=200,
-        num_return_sequences=1,
-        pad_token_id=tokenizer.eos_token_id,
-        do_sample=True,
-        temperature=0.7,
-        top_p=0.9
-    )
-    generated_code = tokenizer.decode(outputs[0], skip_special_tokens=True)
+import os
+import torch
+import json
+from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments, DataCollatorForLanguageModeling, TrainerCallback
+from datasets import Dataset
+import matplotlib.pyplot as plt
+
+# Set Hugging Face token (replace with your actual token)
+os.environ["HF_TOKEN"] = "hf_XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX"  # Replace with your HF_TOKEN
+
+# Download model and tokenizer
+model_name = "Salesforce/codegen-350M-multi"
+local_model_path = "./codegen_model"
+tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=local_model_path)
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float32, cache_dir=local_model_path)
+
+# Set padding token
+tokenizer.pad_token = tokenizer.eos_token
+
+# Move model to CPU
+device = torch.device("cpu")
+model.to(device)
+
+# Load custom dataset from JSONL
+dataset_path = "./custom_dataset.jsonl"
+data = []
+with open(dataset_path, 'r', encoding='utf-8') as f:
+    for line in f:
+        data.append(json.loads(line.strip()))
+dataset = Dataset.from_list(data)
+
+# Tokenize dataset
+def tokenize_function(examples):
+    inputs = [f"{prompt}\n{code}" for prompt, code in zip(examples["prompt"], examples["code"])]
+    return tokenizer(inputs, truncation=True, padding="max_length", max_length=128)
+
+tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=["prompt", "code"])
+
+# Data collator for language modeling
+data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+
+# Define training arguments
+training_args = TrainingArguments(
+    output_dir="./finetuned_codegen",
+    overwrite_output_dir=True,
+    num_train_epochs=5,
+    per_device_train_batch_size=1,
+    gradient_accumulation_steps=4,
+    save_steps=500,
+    save_total_limit=2,
+    logging_steps=10,  # Reduced logging steps for more frequent loss recording
+    learning_rate=5e-5,
+    fp16=False,
+    no_cuda=True,
+    dataloader_pin_memory=False,
+)
+
+# Custom callback to store training loss
+class LossCallback(TrainerCallback):
+    def __init__(self):
+        self.losses = []
+        self.steps = []
+
+    def on_log(self, args, state, control, logs=None, **kwargs):
+        if logs and "loss" in logs:
+            self.losses.append(logs["loss"])
+            self.steps.append(state.global_step)
+
+loss_callback = LossCallback()
+
+# Initialize Trainer
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=tokenized_dataset,
+    data_collator=data_collator,
+    callbacks=[loss_callback],
+)
+
+# Start fine-tuning
+print("Starting fine-tuning...")
+trainer.train()
+
+# Save fine-tuned model
+model.save_pretrained("./finetuned_codegen")
+tokenizer.save_pretrained("./finetuned_codegen")
+
+# Plot training loss
+plt.plot(loss_callback.steps, loss_callback.losses, label="Training Loss")
+plt.xlabel("Steps")
+plt.ylabel("Loss")
+plt.title("Fine-Tuning Loss Curve")
+plt.legend()
+plt.savefig("./finetuned_codegen/loss_plot.png")
+plt.show()
+
+print("Fine-tuning completed. Model saved to ./finetuned_codegen. Loss plot saved to ./finetuned_codegen/loss_plot.png")
+
+# Test fine-tuned model
+print("\nTesting fine-tuned model...")
+prompts = [
+    "Write a Python program to print 'Hello, guys how are you!'"
+]
+
+for prompt in prompts:
+    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=128).to(device)
+    outputs = model.generate(
+        **inputs,
+        max_length=200,
+        num_return_sequences=1,
+        pad_token_id=tokenizer.eos_token_id,
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.9
+    )
+    generated_code = tokenizer.decode(outputs[0], skip_special_tokens=True)
     print(f"Prompt: {prompt}\nGenerated Code:\n{generated_code}\n{'-'*50}")

app.py

@@ -1,39 +1,46 @@
-from flask import Flask, render_template, request
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
-
-app = Flask(__name__)
-
-# Load fine-tuned model and tokenizer
-model_path = "./finetuned_codegen"
-tokenizer = AutoTokenizer.from_pretrained(model_path)
-model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float32)
-
-# Set padding token
-tokenizer.pad_token = tokenizer.eos_token
-
-# Move model to CPU
-device = torch.device("cpu")
-model.to(device)
-
-@app.route("/", methods=["GET", "POST"])
-def index():
-    generated_code = ""
-    prompt = ""
-    if request.method == "POST":
-        prompt = request.form["prompt"]
-        inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=128).to(device)
-        outputs = model.generate(
-            **inputs,
-            max_length=200,
-            num_return_sequences=1,
-            pad_token_id=tokenizer.eos_token_id,
-            do_sample=True,
-            temperature=0.7,
-            top_p=0.9
-        )
-        generated_code = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return render_template("index.html", generated_code=generated_code, prompt=prompt)
-
-if __name__ == "__main__":
+from flask import Flask, render_template, request
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+app = Flask(__name__)
+
+# Load fine-tuned model and tokenizer
+model_path = "./finetuned_codegen"
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float32)
+
+# Set padding token
+tokenizer.pad_token = tokenizer.eos_token
+
+# Move model to CPU
+device = torch.device("cpu")
+model.to(device)
+
+@app.route("/", methods=["GET", "POST"])
+def index():
+    generated_code = ""
+    prompt = ""
+    if request.method == "POST":
+        prompt = request.form["prompt"]
+        inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=128).to(device)
+        outputs = model.generate(
+            **inputs,
+            max_length=200,
+            num_return_sequences=1,
+            pad_token_id=tokenizer.eos_token_id,
+            do_sample=True,
+            temperature=0.2,  # Lower temperature for more precise outputs
+            top_p=0.95,      # Adjusted for better sampling
+            top_k=50,        # Added to focus on top-k tokens
+            no_repeat_ngram_size=3  # Prevent repetitive phrases
+        )
+        generated_code = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        # Clean up output to remove prompt prefix and extra text
+        if generated_code.startswith(prompt):
+            generated_code = generated_code[len(prompt):].strip()
+        # Remove any trailing or redundant text
+        generated_code = generated_code.split("\n")[0].strip() if "\n" in generated_code else generated_code
+    return render_template("index.html", generated_code=generated_code, prompt=prompt)
+
+if __name__ == "__main__":
     app.run(debug=True)

test.py

@@ -4,7 +4,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 # Load fine-tuned model and tokenizer
 model_path = "./finetuned_codegen"
 tokenizer = AutoTokenizer.from_pretrained(model_path)
-model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16)
+model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float32)
 
 # Set padding token
 tokenizer.pad_token = tokenizer.eos_token
@@ -13,7 +13,7 @@ tokenizer.pad_token = tokenizer.eos_token
 device = torch.device("cpu")
 model.to(device)
 
-# Test prompts (including dataset prompts)
+# Test prompts
 prompts = [
     "Write a Python program to print 'Hello, you name or any other thing!'"
 ]