Spaces:
Running
Running
File size: 2,295 Bytes
19a6fbb 8ed49ee 19a6fbb 8ed49ee 19a6fbb 8ed49ee 19a6fbb |
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 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 |
import clip
from datetime import datetime
import faiss
import pandas as pd
import os
from sentence_transformers import SentenceTransformer
import streamlit as st
import torch
import whisper
device = "cuda" if torch.cuda.is_available() else "cpu"
@st.cache_resource
def load_clip_model():
model, preprocess = clip.load("ViT-B/32", device=device)
return model, preprocess
@st.cache_resource
def load_text_embedding_model():
model = SentenceTransformer("all-MiniLM-L6-v2")
return model
@st.cache_resource
def load_whisper_model():
model = whisper.load_model("small")
return model
def load_image_index():
index = faiss.read_index('./vectorstore/image_index.index')
data = pd.read_csv("./vectorstore/image_data.csv")
return index, data
def load_text_index():
index = faiss.read_index('./vectorstore/text_index.index')
data = pd.read_csv("./vectorstore/text_data.csv")
return index, data
def load_audio_index():
index = faiss.read_index('./vectorstore/audio_index.index')
data = pd.read_csv("./vectorstore/audio_data.csv")
return index, data
def cosine_similarity(a, b):
return torch.cosine_similarity(a, b)
def get_local_files(directory: str, extensions: list = None, get_details: bool = False):
files = os.listdir(directory)
if not extensions:
if get_details:
return [{
"file_name": file,
"file_size": os.path.getsize(os.path.join(directory, file)),
"file_created": datetime.fromtimestamp(os.path.getctime(os.path.join(directory, file)))
} for file in files]
else:
return files
else:
if get_details:
filtered_files = []
for file in files:
file_extension = file.split(".")[-1]
if file_extension in extensions:
filtered_files.append({
"file_name": file,
"file_size": os.path.getsize(os.path.join(directory, file)),
"file_created": datetime.fromtimestamp(os.path.getctime(os.path.join(directory, file)))
})
return filtered_files
else:
return [file for file in files if file.endswith(extensions(extensions))]
|