Minor tweaks

lib/api/ingest.py

@@ -1,11 +1,11 @@
 import asyncio
 
-from prisma.models import Datasource
-from fastapi import APIRouter, BackgroundTasks
-from lib.service.flows import create_finetune
+from fastapi import APIRouter
 
 from lib.models.ingest import IngestRequest
+from lib.service.flows import create_finetune
 from lib.utils.prisma import prisma
+from prisma.models import Datasource
 
 router = APIRouter()
 
@@ -15,7 +15,7 @@
     name="ingest",
     description="Ingest data",
 )
-async def ingest(body: IngestRequest, background_tasks: BackgroundTasks):
+async def ingest(body: IngestRequest):
     """Endpoint for ingesting data"""
     datasource = await prisma.datasource.create(data=body.dict())
 

lib/models/ingest.py

@@ -1,6 +1,7 @@
-from pydantic import BaseModel
 from typing import Optional
 
+from pydantic import BaseModel
+
 
 class IngestRequest(BaseModel):
     type: str

lib/service/finetune.py

@@ -1,13 +1,15 @@
+# flake8: noqa
+
 import asyncio
-import httpx
+import json
 import os
 import uuid
-import openai
-import replicate
-import json
-
 from abc import ABC, abstractmethod
 from typing import Dict, List, Union
+
+import httpx
+import openai
+import replicate
 from decouple import config
 from llama_index import Document
 
@@ -109,7 +111,7 @@ class ReplicateFinetuningService(FinetuningService):
     def __init__(
         self,
         nodes: List[Union[Document, None]],
-        num_questions_per_chunk: int = 1,
+        num_questions_per_chunk: int = 10,
         batch_size: int = 10,
         base_model: str = "LLAMA2_7B_CHAT",
     ):
@@ -164,7 +166,7 @@ async def finetune(self, training_file: str) -> Dict:
             version=REPLICATE_MODELS[self.base_model],
             input={
                 "train_data": training_file_url,
-                "num_train_epochs": 3,
+                "num_train_epochs": 6,
             },
             destination="homanp/test",
         )
@@ -212,5 +214,4 @@ async def upload_replicate_dataset(training_file: str) -> str:
             )
 
         serving_url = response_data["serving_url"]
-        print(serving_url)
         return serving_url

lib/service/flows.py

@@ -1,7 +1,7 @@
-import openai
 import json
-
 from typing import List, Union
+
+import openai
 from llama_index import Document
 from prefect import flow, task