Update app.py

app.py

@@ -1,19 +1,27 @@
-# training_space/app.py (Training Space Backend)
+# training_space/app.py (FastAPI Backend)
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 import subprocess
 import os
 import uuid
 from huggingface_hub import HfApi, HfFolder
-import logging
 from fastapi.middleware.cors import CORSMiddleware
+import logging
 
 app = FastAPI()
 
-# Configure CORS
+# Configure Logging
+logging.basicConfig(
+    filename='training.log',
+    filemode='a',
+    format='%(asctime)s - %(levelname)s - %(message)s',
+    level=logging.INFO
+)
+
+# CORS Configuration
 origins = [
-    "https://huggingface.co/spaces/Vishwas1/LLMBuilderPro",  # Replace with your Gradio frontend Space URL
-    "http://localhost",  # If testing locally
+    "https://Vishwas1-LLMBuilderPro.hf.space",  # Replace with your Gradio frontend Space URL
+    "http://localhost",  # For local testing
     "https://web.postman.co",
 ]
 
@@ -25,14 +33,6 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Configure logging
-logging.basicConfig(
-    filename='training.log',
-    filemode='a',
-    format='%(asctime)s - %(levelname)s - %(message)s',
-    level=logging.INFO
-)
-
 # Define the expected payload structure
 class TrainingRequest(BaseModel):
     task: str  # 'generation' or 'classification'
@@ -40,15 +40,7 @@ class TrainingRequest(BaseModel):
     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()
-
+# Root Endpoint
 @app.get("/")
 def read_root():
     return {
@@ -56,10 +48,12 @@ def read_root():
         "instructions": "To train a model, send a POST request to /train with the required parameters."
     }
 
+# Train Endpoint
 @app.post("/train")
 def train_model(request: TrainingRequest):
     try:
         logging.info(f"Received training request for model: {request.model_name}, Task: {request.task}")
+        
         # Create a unique directory for this training session
         session_id = str(uuid.uuid4())
         session_dir = f"./training_sessions/{session_id}"
@@ -76,11 +70,11 @@ def train_model(request: TrainingRequest):
             "--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'])
+            "--num_layers", str(request.model_params.get('num_layers', 12)),
+            "--attention_heads", str(request.model_params.get('attention_heads', 1)),
+            "--hidden_size", str(request.model_params.get('hidden_size', 64)),
+            "--vocab_size", str(request.model_params.get('vocab_size', 30000)),
+            "--sequence_length", str(request.model_params.get('sequence_length', 512))
         ]
         
         # Start the training process as a background task
@@ -106,3 +100,4 @@ def get_status(session_id: str):
         logs = f.read()
     
     return {"session_id": session_id, "logs": logs}
+