8bitnand

class RAGModel:
    def __init__(self, configs) -> None:
        self.configs = configs
        self.device = configs["model"]["device"]
        model_url = configs["model"]["genration_model"]
        # quantization_config = BitsAndBytesConfig(
        #     load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16
        # )

        self.model = AutoModelForCausalLM.from_pretrained(

8bitnand

class SemanticSearch:
    def __init__(
        self, doc_chunks: tuple[list, list], model_path: str, device: str
    ) -> None:

        self.doc_chunks, self.urls = doc_chunks
        self.st = SentenceTransformer(
            model_path,
            device,
        )

8bitnand

def generate(pretrained_pipe_dir):

    pipeline = DDPMPipeline.from_pretrained(pretrained_pipe_dir).to("cuda")
    images = pipeline(
        batch_size=1,
        generator=torch.manual_seed(123),
    ).images

    return images 
  

8bitnand

# source https://huggingface.co/docs/diffusers/tutorials/basic_training

def load_pipline(config):
    
    pipeline = DDPMPipeline.from_pretrained(
        "mrm8488/ddpm-ema-butterflies-128",
        cache_dir="models/pretrained",
    )
    return pipeline.unet