| from fastapi import FastAPI, HTTPException
|
| from pydantic import BaseModel
|
| from fastapi.middleware.cors import CORSMiddleware
|
| import os
|
| import threading
|
| import torch
|
| from app.job_manager import JobManager
|
| from agent.inference import InferenceEngine
|
|
|
| app = FastAPI()
|
|
|
|
|
|
|
| app.add_middleware(
|
| CORSMiddleware,
|
| allow_origins=["*"],
|
| allow_credentials=True,
|
| allow_methods=["*"],
|
| allow_headers=["*"],
|
| )
|
|
|
| class TrainRequest(BaseModel):
|
| data_path: str = None
|
| manual_text: str = None
|
| category: str = "text"
|
| noise_level: float = 0.0
|
|
|
| class AutoTrainRequest(BaseModel):
|
| topic: str = None
|
| url: str = None
|
| category: str = "text"
|
|
|
| class GrammarCheckRequest(BaseModel):
|
| text: str
|
|
|
| class ParallelTrainRequest(BaseModel):
|
| topics: list[str]
|
|
|
| @app.post("/api/train")
|
| async def trigger_training(request: TrainRequest):
|
| if not request.data_path and not request.manual_text:
|
| return {"status": "error", "message": "Either data path or manual text must be provided."}
|
|
|
| path = request.data_path if request.data_path else "manual_input"
|
| if request.data_path and not os.path.exists(path):
|
| return {"status": "error", "message": "Path not found."}
|
|
|
| manager = JobManager()
|
| job = manager.create_job("custom", path)
|
|
|
| def run_train():
|
| try:
|
| print(f"Loading data from {path}...")
|
| print("Starting training...")
|
| import train.train
|
| import importlib
|
| importlib.reload(train.train)
|
|
|
|
|
| train.train.train(
|
| dataset_path=request.data_path,
|
| text_content=request.manual_text,
|
| job=job,
|
| category=request.category,
|
| noise_level=request.noise_level
|
| )
|
|
|
| print("Training Complete. Model saved to sail.pt")
|
| except Exception as e:
|
| print(f"Training Failed: {e}")
|
| job.status = "failed"
|
| job.error = str(e)
|
| job.message = f"Failed: {e}"
|
|
|
| thread = threading.Thread(target=run_train)
|
| job.thread = thread
|
| thread.start()
|
|
|
| return {"status": "success", "message": "Training started.", "job_id": job.id}
|
|
|
| @app.post("/api/train/auto")
|
| async def trigger_auto_training(request: AutoTrainRequest):
|
| if not request.topic and not request.url:
|
| return {"status": "error", "message": "Either topic or url must be provided."}
|
|
|
| manager = JobManager()
|
| details = request.topic if request.topic else request.url
|
| job = manager.create_job("auto", details)
|
|
|
| def run_auto_train():
|
| try:
|
| from agent.data_collector import DataCollector
|
| collector = DataCollector()
|
|
|
| text_data = ""
|
| job.message = "Collecting data..."
|
| if request.topic:
|
| print(f"Collecting data for topic: {request.topic}")
|
| text_data = collector.search_and_collect(request.topic, category=request.category)
|
| elif request.url:
|
| print(f"Collecting data from URL: {request.url}")
|
| text_data = collector.collect_from_url(request.url)
|
|
|
| if not text_data:
|
| print("No data collected.")
|
| job.status = "failed"
|
| job.error = "No data collected"
|
| job.message = "Failed: No data collected"
|
| return
|
|
|
| print("Starting automatic training...")
|
| import train.train
|
| import importlib
|
| importlib.reload(train.train)
|
|
|
| train.train.train(text_content=text_data, job=job, category=request.category)
|
| print("Auto Training Complete.")
|
|
|
| except Exception as e:
|
| print(f"Auto Training Failed: {e}")
|
| job.status = "failed"
|
| job.error = str(e)
|
| job.message = f"Failed: {e}"
|
|
|
| thread = threading.Thread(target=run_auto_train)
|
| job.thread = thread
|
| thread.start()
|
|
|
| return {"status": "success", "message": "Automatic training started.", "job_id": job.id}
|
|
|
| @app.post("/api/train/wordnet")
|
| def trigger_wordnet_training(request: AutoTrainRequest):
|
| """Generates grammar instructions from WordNet and trains."""
|
| manager = JobManager()
|
| job = manager.create_job("wordnet", f"Grammar Training (WordNet)")
|
|
|
| def run_wordnet_train():
|
| try:
|
| from train.wordnet_loader import WordNetLoader
|
| import json
|
|
|
| loader = WordNetLoader()
|
| job.message = "Loading WordNet (slower first time)..."
|
|
|
|
|
| print("Generating WordNet dataset...")
|
| data = loader.generate_instruction_dataset(num_samples=2000)
|
|
|
| if not data:
|
| job.status = "failed"
|
| job.error = "No WordNet data generated"
|
| return
|
|
|
|
|
| dataset_path = "wordnet_grammar.json"
|
| with open(dataset_path, 'w', encoding='utf-8') as f:
|
| json.dump(data, f, indent=2)
|
|
|
| print(f"Saved {len(data)} grammar instructions to {dataset_path}")
|
|
|
| print("Starting grammar training...")
|
| import train.train
|
| import importlib
|
| importlib.reload(train.train)
|
|
|
| train.train.train(dataset_path=dataset_path, job=job, category="grammar")
|
| print("Grammar Training Complete")
|
|
|
| except Exception as e:
|
| print(f"Grammar Training Failed: {e}")
|
| job.status = "failed"
|
| job.error = str(e)
|
| job.message = f"Failed: {e}"
|
|
|
| thread = threading.Thread(target=run_wordnet_train)
|
| job.thread = thread
|
| thread.start()
|
|
|
| return {"status": "started", "job_id": job.id, "message": "Grammar training started."}
|
|
|
| @app.post("/api/train/instruction")
|
| def trigger_instruction_training(request: AutoTrainRequest):
|
| """Generates Q&A pairs and trains the model on them."""
|
| manager = JobManager()
|
| job = manager.create_job("instruction", f"Category: {request.category}")
|
|
|
| def run_instruction_train():
|
| try:
|
| from agent.data_collector import DataCollector
|
| import json
|
|
|
| collector = DataCollector()
|
| job.message = f"Generating {request.category} instructions..."
|
| print(f"Generating instructions for {request.category}...")
|
|
|
| instructions = collector.collect_instructions(request.category)
|
|
|
| if not instructions:
|
| job.status = "failed"
|
| job.error = "No instructions generated"
|
| return
|
|
|
| dataset_path = "instructions.json"
|
| with open(dataset_path, 'w', encoding='utf-8') as f:
|
| json.dump(instructions, f, indent=2)
|
|
|
| print(f"Saved {len(instructions)} instructions to {dataset_path}")
|
|
|
| print("Starting instruction training...")
|
| import train.train
|
| import importlib
|
| importlib.reload(train.train)
|
|
|
| train.train.train(dataset_path=dataset_path, job=job, category=request.category)
|
| print("Instruction Training Complete")
|
|
|
| except Exception as e:
|
| print(f"Instruction Training Failed: {e}")
|
| job.status = "failed"
|
| job.error = str(e)
|
| job.message = f"Failed: {e}"
|
|
|
| thread = threading.Thread(target=run_instruction_train)
|
| job.thread = thread
|
| thread.start()
|
|
|
| return {"status": "started", "job_id": job.id, "message": "Instruction training started."}
|
|
|
| @app.post("/api/train/parallel")
|
| def trigger_parallel_training(request: ParallelTrainRequest):
|
| """Hits the new Multi-Stream single-GPU orchestrator to train multiple topics."""
|
| if not request.topics or len(request.topics) == 0:
|
| return {"status": "error", "message": "List of topics required."}
|
|
|
| manager = JobManager()
|
|
|
| def run_parallel_train():
|
| from agent.data_collector import DataCollector
|
| import train.train
|
| import importlib
|
| importlib.reload(train.train)
|
|
|
| collector = DataCollector()
|
| tasks = []
|
|
|
| for topic in request.topics:
|
| job = manager.create_job("parallel", f"Parallel: {topic}")
|
| job.message = f"Collecting {topic}..."
|
| print(f"Collecting data for {topic}...")
|
|
|
| data = collector.search_and_collect(topic, category="custom")
|
| if data:
|
| tasks.append({
|
| "text_content": data,
|
| "category": f"parallel_{topic}",
|
| "job": job
|
| })
|
| else:
|
| job.status = "failed"
|
| job.error = "No data found."
|
| job.message = "Failed: No Data."
|
|
|
| if len(tasks) > 0:
|
| print(f"Dispatching {len(tasks)} tasks to concurrent CUDA Streams...")
|
| train.train.train_parallel(tasks)
|
| else:
|
| print("No parallel tasks gathered data.")
|
|
|
| thread = threading.Thread(target=run_parallel_train)
|
| thread.start()
|
|
|
| return {"status": "started", "message": f"Parallel multi-stream training started for {len(request.topics)} topics."}
|
|
|
| @app.get("/api/jobs")
|
| async def get_jobs():
|
| manager = JobManager()
|
| return {"jobs": manager.get_all_jobs()}
|
|
|
| @app.post("/api/jobs/{job_id}/{action}")
|
| async def control_job(job_id: str, action: str):
|
| manager = JobManager()
|
| success = False
|
| if action == "pause":
|
| success = manager.pause_job(job_id)
|
| elif action == "resume":
|
| success = manager.resume_job(job_id)
|
| elif action == "terminate":
|
| success = manager.terminate_job(job_id)
|
|
|
| if success:
|
| return {"status": "success", "message": f"Job {action}d"}
|
| else:
|
| return {"status": "error", "message": f"Failed to {action} job"}
|
|
|
| @app.post("/api/reset")
|
| async def reset_brain():
|
| """Deletes the model file to start fresh."""
|
| import glob
|
| try:
|
| if os.path.exists("sail.pt"):
|
| os.remove("sail.pt")
|
|
|
| for f in glob.glob("sail_*.pt"):
|
| os.remove(f)
|
|
|
| print("Brain reset successfully.")
|
| return {"status": "success", "message": "Brain reset. Please reload the page."}
|
| except Exception as e:
|
| return {"status": "error", "message": str(e)}
|
|
|
| @app.get("/api/browse")
|
| async def browse_folder():
|
| """Opens a native folder picker dialog on the server and returns the path."""
|
| import subprocess
|
| try:
|
|
|
| cmd = [
|
| "python", "-c",
|
| "import tkinter.filedialog as fd; import tkinter; root = tkinter.Tk(); root.withdraw(); root.attributes('-topmost', True); print(fd.askdirectory()); root.destroy()"
|
| ]
|
| result = subprocess.run(cmd, capture_output=True, text=True)
|
| path = result.stdout.strip()
|
| return {"path": path}
|
| except Exception as e:
|
| return {"path": "", "error": str(e)}
|
|
|
| @app.get("/api/status")
|
| async def status():
|
| manager = JobManager()
|
| active_jobs = [j for j in manager.jobs.values() if j.status == "running"]
|
| return {
|
| "is_training": len(active_jobs) > 0,
|
| "model_exists": os.path.exists("sail.pt"),
|
| "active_jobs_count": len(active_jobs)
|
| }
|
|
|
