1

app.py

@@ -10,34 +10,37 @@ model = GPT2LMHeadModel.from_pretrained('gpt2')
 tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")
 tokenizer.pad_token = tokenizer.eos_token
 
-temp_folder = 'temp'
-os.makedirs(temp_folder, exist_ok=True)
+# temp_folder = 'temp'
+# os.makedirs(temp_folder, exist_ok=True)
+logit = {}
 json_file = 'index.json'
 with open(json_file, 'r') as file:
     data = json.load(file)
 for key, value in data.items():
     text_description = value['text_description']
-    inputs = tokenizer(text_description, return_tensors="pt", padding="max_length", max_length=128, truncation=True)
+    inputs = tokenizer(text_description, return_tensors="pt", padding="max_length", max_length=32, truncation=True)
     outputs = model(**inputs, labels=inputs["input_ids"])
     logits = outputs.logits
-    torch.save(logits, os.path.join(temp_folder, f"{key}.pt"))
+    logit[key] = logits
+    # torch.save(logits, os.path.join(temp_folder, f"{key}.pt"))
 
 
 def search_index(query):
-    inputs = tokenizer(query, return_tensors="pt", padding="max_length", max_length=128, truncation=True)
+    inputs = tokenizer(query, return_tensors="pt", padding="max_length", max_length=32, truncation=True)
     outputs = model(**inputs, labels=inputs["input_ids"])
 
     max_similarity = float('-inf')
     max_similarity_uuid = None
-    for file in os.listdir(temp_folder):
-        uuid = file.split('.')[0]
-        logits = torch.load(os.path.join(temp_folder, file))
+    # for file in os.listdir(temp_folder):
+    #     uuid = file.split('.')[0]
+    #     logits = torch.load(os.path.join(temp_folder, file))
+    for uuid, logits in logit.items():
         similarity = (outputs.logits * logits).sum()
         if similarity > max_similarity:
             max_similarity = similarity
             max_similarity_uuid = uuid
 
-    gr.logger.info(f"Query: {query}")
+    gr.Info(f"Max similarity: {max_similarity}")
     return max_similarity_uuid