modified: app.py

app.py

@@ -3,14 +3,38 @@ import os
 from huggingface_hub import hf_hub_download
 from pathlib import Path
 from transformers import GPT2Config, GPT2LMHeadModel, GPT2Tokenizer
+import json
 
-config_class, model_class, tokenizer_class = GPT2Config, GPT2LMHeadModel, GPT2Tokenizer
-model = model_class.from_pretrained('gpt2')
+model = GPT2LMHeadModel.from_pretrained('gpt2')
 tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
 
+logits_dict = {}
+
+json_file = 'index.json'
+with open(json_file, 'r') as file:
+    data = json.load(file)
+for item in data:
+    uuid = item['uuid']
+    text_description = item['text_description']
+    inputs = tokenizer(text_description, return_tensors="pt", padding=True, truncation=True)
+    outputs = model(**inputs, labels=inputs["input_ids"])
+    logits = outputs.logits
+    logits_dict[uuid] = logits
+
 
 def search_index(query):
-    return "example_uuid"
+    inputs = tokenizer(query, return_tensors="pt", padding=True, truncation=True)
+    outputs = model(**inputs, labels=inputs["input_ids"])
+
+    max_similarity = float('-inf')
+    max_similarity_uuid = None
+    for uuid, logits in logits_dict.items():
+        similarity = (outputs.logits * logits).sum()
+        if similarity > max_similarity:
+            max_similarity = similarity
+            max_similarity_uuid = uuid
+
+    return max_similarity_uuid
 
 
 def download_video(uuid):