subashdvorak commited on
Commit
509759d
·
verified ·
1 Parent(s): 5188e99

Update shared_resources.py

Browse files
Files changed (1) hide show
  1. shared_resources.py +24 -24
shared_resources.py CHANGED
@@ -1,24 +1,24 @@
1
- # shared_resources.py
2
- import torch
3
- # from transformers import AutoTokenizer, AutoModelForCausalLM
4
- from sentence_transformers import SentenceTransformer
5
- from datasets import load_dataset
6
- # from transformers import AutoProcessor, MusicgenForConditionalGeneration
7
- import re
8
-
9
- class SharedResources:
10
- def __init__(self):
11
- # Set the device
12
- self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
13
-
14
-
15
- # Load SentenceTransformer
16
- self.sentence_transformer = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")
17
-
18
- # Load the dataset
19
- self.dataset = load_dataset("subashdvorak/tiktok-story-data1", revision="embedded")
20
- self.data = self.dataset["train"]
21
- self.data = self.data.add_faiss_index("embeddings")
22
-
23
- # Create a single instance of SharedResources
24
- shared_resources = SharedResources()
 
1
+ # shared_resources.py
2
+ import torch
3
+ # from transformers import AutoTokenizer, AutoModelForCausalLM
4
+ from sentence_transformers import SentenceTransformer
5
+ from datasets import load_dataset
6
+ # from transformers import AutoProcessor, MusicgenForConditionalGeneration
7
+ import re
8
+
9
+ class SharedResources:
10
+ def __init__(self):
11
+ # Set the device
12
+ self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
13
+
14
+
15
+ # Load SentenceTransformer
16
+ self.sentence_transformer = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")
17
+
18
+ # Load the dataset
19
+ self.dataset = load_dataset("subashdvorak/tiktok-story-data3", revision="embedded")
20
+ self.data = self.dataset["train"]
21
+ self.data = self.data.add_faiss_index("embeddings")
22
+
23
+ # Create a single instance of SharedResources
24
+ shared_resources = SharedResources()