Spaces:
Build error
Build error
| import nmslib | |
| import streamlit as st | |
| from transformers import CLIPProcessor, AutoTokenizer, ViTFeatureExtractor | |
| import numpy as np | |
| from koclip import FlaxHybridCLIP | |
| def load_index(img_file): | |
| filenames, embeddings = [], [] | |
| lines = open(img_file, "r") | |
| for line in lines: | |
| cols = line.strip().split('\t') | |
| filename = cols[0] | |
| embedding = np.array([float(x) for x in cols[1].split(',')]) | |
| filenames.append(filename) | |
| embeddings.append(embedding) | |
| embeddings = np.array(embeddings) | |
| index = nmslib.init(method='hnsw', space='cosinesimil') | |
| index.addDataPointBatch(embeddings) | |
| index.createIndex({'post': 2}, print_progress=True) | |
| return filenames, index | |
| def load_model(model_name="koclip/koclip"): | |
| assert model_name in {"koclip/koclip", "koclip/koclip-large"} | |
| model = FlaxHybridCLIP.from_pretrained(model_name) | |
| processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32") | |
| processor.tokenizer = AutoTokenizer.from_pretrained("klue/roberta-large") | |
| if model_name == "koclip/koclip-large": | |
| processor.feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-large-patch16-224") | |
| return model, processor | |
| def load_model_v2(model_name="koclip/koclip"): | |
| model = FlaxHybridCLIP.from_pretrained(model_name) | |
| processor = CLIPProcessor.from_pretrained(model_name) | |
| return model, processor | |