1  from os import system
  2  from dataclasses import dataclass
  3  from sqlite3 import connect, Error as SQLError
  4  from ollama import create, RequestError, ResponseError
  5  from loguru import logger
  6  
  7  DB_PATH: str = "db/models.db"
  8  
  9  name: str = input("Choose a name for your model: ")
 10  select: str = input("Select a model to train: ")
 11  prompt: str = input("Enter a prompt to train the model with: ")
 12  backup: str = input("Would you like to backup this model? (y/n/restore/view): ")
 13  
 14  system("clear")
 15  
 16  modelfile: str = f"""
 17  FROM {select}
 18  SYSTEM {prompt}
 19  """
 20  
 21  
 22  class Database:
 23      @staticmethod
 24      def create_backup(path: str) -> None:
 25          try:
 26              conn = connect(path)
 27              c = conn.cursor()
 28              c.execute("CREATE TABLE IF NOT EXISTS models (name TEXT, modelfile TEXT)")
 29              c.execute("INSERT INTO models VALUES (?, ?)", (name, modelfile))
 30              conn.commit()
 31          except SQLError as e:
 32              logger.error(e)
 33          finally:
 34              conn.close()
 35  
 36      @staticmethod
 37      def restore_backup(path: str) -> None:
 38          try:
 39              conn = connect(path)
 40              c = conn.cursor()
 41              c.execute("SELECT * FROM models")
 42              models = c.fetchall()
 43              for model in models:
 44                  create(model=model[0], modelfile=model[1])
 45                  logger.info(f"Restored Model: {model[0]}")
 46          except SQLError as e:
 47              logger.error(e)
 48          finally:
 49              conn.close()
 50      
 51      @staticmethod
 52      def lookup(path: str, name: str) -> None:
 53          try:
 54              conn = connect(path)
 55              c = conn.cursor()
 56              c.execute("SELECT modelfile FROM models WHERE name=?", (name,))
 57              modelfile = c.fetchone()
 58              logger.info(f"Model: {name}\n {modelfile[0]}")
 59          except SQLError as e:
 60              logger.error(e)
 61          finally:
 62              conn.close()
 63  
 64  
 65  class Factory:
 66      @staticmethod
 67      def generate() -> None:
 68          try:
 69              model: Model = Model(name, modelfile)
 70              create(model=model.name, modelfile=model.modelfile)
 71              logger.info(f"Generated Model: {model.name}")
 72          except (RuntimeError, RequestError, ResponseError) as e:
 73              logger.error(e)
 74          finally:
 75              match backup:
 76                  case "y":
 77                      Database.create_backup(DB_PATH)
 78                      logger.info(f"Backup created for Model: {name}")
 79                  case "n":
 80                      logger.info("No backup created")
 81                      exit(0)
 82                  case "restore":
 83                      logger.info("Restoring backup...")
 84                      Database.restore_backup(DB_PATH)
 85                      exit(0)
 86                  case "view":
 87                      query = input("Enter the name of the backup you'd like to view: ")
 88                      system("clear")
 89                      Database.lookup(DB_PATH, query)
 90                      exit(0)
 91                  case _:
 92                      logger.error("No backup created. Invalid input, exiting...")
 93                      exit(1)
 94  
 95  
 96  @dataclass
 97  class Model(Factory):
 98      name: str
 99      modelfile: str
100      
101      
102  if __name__ == "__main__":
103      Model.generate()