Create app.py

app.py

@@ -0,0 +1,59 @@
+# training_space/app.py (Training Space Backend)
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+import subprocess
+import os
+import uuid
+from transformers import HfApi, HfFolder
+
+app = FastAPI()
+
+# Define the expected payload structure
+class TrainingRequest(BaseModel):
+    task: str  # 'generation' or 'classification'
+    model_params: dict
+    model_name: str
+    dataset_content: str  # The actual content of the dataset
+
+# Ensure Hugging Face API token is set as an environment variable
+HF_API_TOKEN = os.getenv("HF_API_TOKEN")
+if not HF_API_TOKEN:
+    raise ValueError("HF_API_TOKEN environment variable not set.")
+
+# Save the token
+HfFolder.save_token(HF_API_TOKEN)
+api = HfApi()
+
+@app.post("/train")
+def train_model(request: TrainingRequest):
+    try:
+        # Create a unique directory for this training session
+        session_id = str(uuid.uuid4())
+        session_dir = f"./training_sessions/{session_id}"
+        os.makedirs(session_dir, exist_ok=True)
+        
+        # Save the dataset content to a file
+        dataset_path = os.path.join(session_dir, "dataset.txt")
+        with open(dataset_path, "w", encoding="utf-8") as f:
+            f.write(request.dataset_content)
+        
+        # Prepare the command to run the training script
+        cmd = [
+            "python", "train_model.py",
+            "--task", request.task,
+            "--model_name", request.model_name,
+            "--dataset", dataset_path,
+            "--num_layers", str(request.model_params['num_layers']),
+            "--attention_heads", str(request.model_params['attention_heads']),
+            "--hidden_size", str(request.model_params['hidden_size']),
+            "--vocab_size", str(request.model_params['vocab_size']),
+            "--sequence_length", str(request.model_params['sequence_length'])
+        ]
+        
+        # Start the training process as a background task
+        subprocess.Popen(cmd, cwd=session_dir)
+        
+        return {"status": "Training started", "session_id": session_id}
+    
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))