tl-playground / model.py
Andreas Sünder
Add files from previous repo
fce98ea
raw
history blame
516 Bytes
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)