Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from transformers import GPT2Tokenizer, GPT2LMHeadModel, Trainer, TrainingArguments, DataCollatorForLanguageModeling | |
| from datasets import Dataset | |
| import os | |
| # Load the GPT-2 model and tokenizer | |
| tokenizer = GPT2Tokenizer.from_pretrained("gpt2") | |
| model = GPT2LMHeadModel.from_pretrained("gpt2") | |
| # Add a padding token to the tokenizer | |
| tokenizer.pad_token = tokenizer.eos_token | |
| # Load and format your custom dataset for Q&A training | |
| def load_custom_dataset(file_path): | |
| with open(file_path, "r", encoding="utf-8") as file: | |
| lines = file.readlines() | |
| questions, answers = [], [] | |
| for line in lines: | |
| if line.startswith("Q:"): | |
| questions.append(line.strip()) | |
| elif line.startswith("A:"): | |
| answers.append(line.strip()) | |
| print(f"Loaded {len(questions)} questions and {len(answers)} answers.") | |
| return Dataset.from_dict({"questions": questions, "answers": answers}) | |
| dataset_file = "stroon_tutor.txt" # Make sure your file is formatted correctly | |
| if os.path.exists(dataset_file): | |
| dataset = load_custom_dataset(dataset_file) | |
| else: | |
| raise FileNotFoundError(f"File {dataset_file} not found.") | |
| # Tokenize the dataset with distinct question and answer tokens | |
| def tokenize_qa(examples): | |
| input_texts = [q for q in examples["questions"]] | |
| target_texts = [a for a in examples["answers"]] | |
| model_inputs = tokenizer(input_texts, max_length=128, truncation=True, padding="max_length") | |
| labels = tokenizer(target_texts, max_length=128, truncation=True, padding="max_length") | |
| model_inputs["labels"] = labels["input_ids"] | |
| return model_inputs | |
| tokenized_dataset = dataset.map(tokenize_qa, batched=True) | |
| # Check the size of the tokenized dataset | |
| print(f"Tokenized dataset has {len(tokenized_dataset)} entries.") # Print dataset size after tokenization | |
| # Set up training arguments | |
| training_args = TrainingArguments( | |
| output_dir="./results", | |
| evaluation_strategy="no", | |
| learning_rate=1e-5, | |
| per_device_train_batch_size=2, | |
| num_train_epochs=3, | |
| weight_decay=0.01, | |
| ) | |
| # Define the Trainer | |
| data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False) | |
| trainer = Trainer( | |
| model=model, | |
| args=training_args, | |
| train_dataset=tokenized_dataset, | |
| data_collator=data_collator | |
| ) | |
| # Fine-tune the model | |
| trainer.train() | |
| # Function to generate responses from the fine-tuned model | |
| def generate_response(prompt): | |
| inputs = tokenizer(prompt, return_tensors="pt") | |
| outputs = model.generate( | |
| **inputs, | |
| max_length=50, | |
| num_beams=5, | |
| no_repeat_ngram_size=2, | |
| do_sample=True, | |
| top_k=50, | |
| top_p=0.95 | |
| ) | |
| response = tokenizer.decode(outputs[0], skip_special_tokens=True) | |
| return response | |
| # Set up Gradio interface for interaction | |
| iface = gr.Interface(fn=generate_response, inputs="text", outputs="text", title="GPT-2 Fine-Tuned Assistant") | |
| iface.launch() | |