Update app.py

app.py

@@ -10,6 +10,8 @@ __checkpoint = "gpt2"
 __tokenizer = GPT2Tokenizer.from_pretrained(__checkpoint)
 __model = GPT2LMHeadModel.from_pretrained(__checkpoint)
 __model_output_path = "paavansundar/Medical_QNA_GPT2"
+# Create a Data collator object
+__data_collator = DataCollatorForLanguageModeling(tokenizer=__tokenizer, mlm=False, return_tensors="pt")
 
 #prepare data
 def prepareData():
@@ -35,14 +37,44 @@ def prepareData():
     with open("val.txt", "w") as f:
         f.writelines(line+'\n' for line in val_seq)
     
-# Create a Data collator object
-data_collator = DataCollatorForLanguageModeling(tokenizer=__tokenizer, mlm=False, return_tensors="pt")
+
+
+def fine_tune_gpt():
+    
+    model_output_path = "gpt_model"
+    train_dataset = TextDataset(tokenizer=tokenizer, file_path="train.txt", block_size=128)
+    val_dataset = TextDataset(tokenizer=tokenizer, file_path="val.txt", block_size=128)
+    training_args = TrainingArguments(
+        output_dir = model_output_path,
+        overwrite_output_dir = True,
+        per_device_train_batch_size = 2, # try with 2
+        per_device_eval_batch_size = 2,  #  try with 2
+        num_train_epochs = 0.01,
+        save_steps = 1_000,
+        save_total_limit = 2,
+        logging_dir = './logs',
+        )
+
+# Train the model
+    trainer = Trainer(
+        model = __model,
+        args = training_args,
+        data_collator = __data_collator,
+        train_dataset = train_dataset,
+        eval_dataset = val_dataset,
+    )
+
+    trainer.train()
+
+    # Save the model
+    trainer.save_model(model_output_path)
+
+# Save the tokenizer
+tokenizer.save_pretrained(model_output_path)
 def queryGPT(question):
     return generate_response(__model, __tokenizer, question)
 
 def generate_response(model,tokenizer, prompt, max_length=200):
-    train_dataset = TextDataset(tokenizer=tokenizer, file_path="train.txt", block_size=128)
-    val_dataset = TextDataset(tokenizer=tokenizer, file_path="val.txt", block_size=128)
     input_ids = tokenizer.encode(prompt, return_tensors="pt")      # 'pt' for returning pytorch tensor
 
     # Create the attention mask and pad token id
@@ -67,4 +99,5 @@ with gr.Blocks() as demo:
     btn.click(queryGPT, inputs=[txt_input], outputs=[txt_output])
 if __name__ == "__main__":
     prepareData()
+    fine_tune_gpt()
     demo.launch()