Created app.py

app.py

@@ -0,0 +1,77 @@
+import gradio as gr
+from datasets import load_dataset, Dataset
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments
+import torch
+
+# Load and preprocess dataset
+def load_and_preprocess_data():
+    # Load the dataset from the uploaded JSON file
+    dataset = load_dataset('json', data_files='data/your_dataset.json')  # Update path to your JSON file
+
+    # Extract "prompt" and "chosen"/"rejected" as input-output pairs
+    def process_row(example):
+        return {
+            'prompt': example['row']['prompt'],
+            'response': example['row']['chosen'],  # You can change this to 'rejected' if needed
+        }
+    
+    dataset = dataset.map(process_row)
+
+    # Load the tokenizer (You can choose any model, here we use BERT)
+    tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')
+
+    # Preprocess the dataset: tokenize inputs and labels
+    def preprocess_function(examples):
+        return tokenizer(examples['prompt'], truncation=True, padding='max_length', max_length=512)
+
+    # Apply preprocessing to the dataset
+    dataset = dataset.map(preprocess_function, batched=True)
+
+    return dataset, tokenizer
+
+# Define and train the model
+def train_model(dataset):
+    # Load the model (BERT-based for sequence classification)
+    model = AutoModelForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=2)  # Change num_labels if needed
+
+    # Define training arguments
+    training_args = TrainingArguments(
+        output_dir='./results',          # Output directory for model checkpoints
+        evaluation_strategy="epoch",     # Evaluation at the end of each epoch
+        learning_rate=2e-5,              # Learning rate
+        per_device_train_batch_size=8,   # Batch size for training
+        per_device_eval_batch_size=8,    # Batch size for evaluation
+        num_train_epochs=3,              # Number of training epochs
+        weight_decay=0.01,               # Strength of weight decay
+    )
+
+    # Initialize the Trainer
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=dataset['train'],
+        eval_dataset=dataset['test'],  # If you have a test split, otherwise just use train for both
+    )
+
+    # Train the model
+    trainer.train()
+
+    return model
+
+# Gradio Interface function to predict
+def predict(text):
+    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
+    outputs = model(**inputs)
+    prediction = torch.argmax(outputs.logits, dim=-1).item()  # Get the class with the highest score
+    return f"Prediction: {prediction}"
+
+# Main script to run training
+if __name__ == "__main__":
+    # Load and preprocess the data
+    dataset, tokenizer = load_and_preprocess_data()
+
+    # Train the model
+    model = train_model(dataset)
+
+    # Create a Gradio interface for prediction
+    gr.Interface(fn=predict, inputs=gr.Textbox(), outputs="text").launch()