from abc import ABC, abstractmethod
import replicate

class GenericModel:
  def __init__(self, name: str):
    self.name = name

  @abstractmethod
  def load(self):
    pass

  @abstractmethod
  def generate(self, prompt: str):
    pass

class ReplicateModel(GenericModel):
  def __init__(self, name: str):
    super().__init__(name)

  def load(self):
    return self

  def generate(self, prompt: str):
    iterator = replicate.run(self.name, input={'prompt': prompt})
    return ''.join(item for item in iterator)