VectorDB-and-LLM/setup_req.py at main · engineerjkk/VectorDB-and-LLM · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
from pinecone import Pinecone
from image_utils import fetch_clip
from splade.splade.models.transformer_rep import Splade
from transformers import AutoTokenizer
from transformers import YolosFeatureExtractor, YolosForObjectDetection
import pandas as pd
import json
from config import MY_PINECONE_API_KEY
import warnings
warnings.filterwarnings("ignore")

# setup

def setup():
    #### pinecone ####
    pc = Pinecone(api_key=MY_PINECONE_API_KEY)

    index = pc.Index("engineerjkk")
    index.describe_index_stats()

    #### CLIP model ####
    clip_model, clip_processor, clip_tokenizer = fetch_clip(model_name="patrickjohncyh/fashion-clip")


    #### SPLADE model ####
    splade_model_id = 'naver/splade-cocondenser-ensembledistil'

    splade_model = Splade(splade_model_id, agg='max')
    splade_model.to('cpu')
    splade_model.eval()

    splade_tokenizer = AutoTokenizer.from_pretrained(splade_model_id)

    #### local DB ####
    local_db = pd.read_csv("local_db.csv")
    local_db['values'] = local_db['values'].apply(json.loads)

    #### YOLO model ####
    MODEL_NAME = "valentinafeve/yolos-fashionpedia"

    yolo_feature_extractor = YolosFeatureExtractor.from_pretrained('hustvl/yolos-small')
    yolo_model = YolosForObjectDetection.from_pretrained(MODEL_NAME)

    return pc, index, clip_model, clip_processor, clip_tokenizer, splade_model, splade_tokenizer, local_db, yolo_feature_extractor, yolo_model