Update app.py

app.py

@@ -1,27 +1,45 @@
-from transformers import AutoModelForSeq2SeqLM, GPT2Tokenizer
-import gradio as gr
-# Load your custom tokenizer
-custom_tokenizer = GPT2Tokenizer.from_pretrained("athirdpath/DPO_Pairs-Roleplay-Alpaca-NSFW")
-
-# Load the pre-trained model
-model = AutoModelForSeq2SeqLM.from_pretrained("facebook/bart-large-cnn", from_tf=True)
-
-@spaces.GPU
-def generate(prompt):
-    # Tokenize your input using the custom tokenizer
-    inputs = custom_tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True)
-
-    # Generate output using the pre-trained model
-    generated_ids = model.generate(inputs["input_ids"], max_new_tokens=150)
-    
-    # Decode the generated output
-    output = custom_tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
-
-    return output
-
-input_component = gr.Textbox(label="Input a persona, e.g. photographer", value="photographer")
-output_component = gr.Textbox(label="Prompt")
-examples = [["Hello!"], ["Are you a chatbot?"]]
-description = "This app generates VeoAI Beta Prompts, based on a model trained on a custom dataset."
-gr.Interface(generate, inputs=input_component, outputs=output_component, examples=examples, title="VeoAI",
-             description=description).launch()
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, Text2TextGenerationPipeline
+from transformers import Trainer, TrainingArguments
+from datasets import load_dataset
+
+# Load your custom dataset
+dataset = load_dataset("athirdpath/DPO_Pairs-Roleplay-Alpaca-NSFW")
+
+# Load the pre-trained model and tokenizer
+model_name = "facebook/bart-large-cnn"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+
+# Define training arguments
+training_args = TrainingArguments(
+    output_dir="./veo_small",
+    num_train_epochs=3,
+    per_device_train_batch_size=4,
+    save_steps=10_000,
+    save_total_limit=2,
+)
+
+# Define Trainer
+trainer = Trainer(
+    model=model,
+    tokenizer=tokenizer,
+    train_dataset=dataset["train"],
+    args=training_args,
+)
+
+# Train the model
+trainer.train()
+
+# Save the fine-tuned model
+model.save_pretrained("./veo_small")
+tokenizer.save_pretrained("./veo_small")
+
+# Use the fine-tuned model for generation
+custom_model = AutoModelForSeq2SeqLM.from_pretrained("./veo_small")
+custom_tokenizer = AutoTokenizer.from_pretrained("./veo_small")
+generator = Text2TextGenerationPipeline(model=custom_model, tokenizer=custom_tokenizer)
+
+# Example of model usage
+prompt = "Hi"
+output = generator(prompt, max_length=150, num_return_sequences=1)
+print(output[0]["generated_text"])