1

app.py

@@ -4,13 +4,14 @@ from huggingface_hub import hf_hub_download
 from pathlib import Path
 from transformers import GPT2Config, GPT2LMHeadModel, GPT2TokenizerFast
 import json
+import torch
 
 model = GPT2LMHeadModel.from_pretrained('gpt2')
 tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")
 tokenizer.pad_token = tokenizer.eos_token
 
-logits_dict = {}
-
+temp_folder = 'temp'
+os.mkdir(temp_folder, exist_ok=True)
 json_file = 'index.json'
 with open(json_file, 'r') as file:
     data = json.load(file)
@@ -19,7 +20,7 @@ for key, value in data.items():
     inputs = tokenizer(text_description, return_tensors="pt", padding="max_length", max_length=128, truncation=True)
     outputs = model(**inputs, labels=inputs["input_ids"])
     logits = outputs.logits
-    logits_dict[key] = logits
+    torch.save(logits, os.path.join(temp_folder, f"{key}.pt"))
 
 
 def search_index(query):
@@ -28,7 +29,9 @@ def search_index(query):
 
     max_similarity = float('-inf')
     max_similarity_uuid = None
-    for uuid, logits in logits_dict.items():
+    for file in os.listdir(temp_folder):
+        uuid = file.split('.')[0]
+        logits = torch.load(os.path.join(temp_folder, file))
         similarity = (outputs.logits * logits).sum()
         if similarity > max_similarity:
             max_similarity = similarity