Update inference.py

inference.py

@@ -1,33 +1,33 @@
-import torch
-from model import LVL
-from transformers import RobertaTokenizer
-from PIL import Image
-from torchvision import transforms
-
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# Load model
-model = LVL()
-model.load_state_dict(torch.load("pytorch_model.bin", map_location=device))
-model.to(device)
-model.eval()
-
-# Load tokenizer
-tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
-
-# Image transform
-transform = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor()
-])
-
-
-def predict(image_path, text):
-    image = transform(Image.open(image_path).convert("RGB")).unsqueeze(0).to(device)
-    tokens = tokenizer(text, return_tensors="pt", padding=True, truncation=True).to(device)
-
-    with torch.no_grad():
-        img_feat, txt_feat = model(image, tokens["input_ids"], tokens["attention_mask"])
-        similarity = torch.matmul(img_feat, txt_feat.T).squeeze()
-
-    return similarity.item()
+import torch
+from model import LVL
+from transformers import RobertaTokenizer
+from PIL import Image
+from torchvision import transforms
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Load model
+model = LVL()
+model.load_state_dict(torch.load("scold.pth", map_location=device))
+model.to(device)
+model.eval()
+
+# Load tokenizer
+tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
+
+# Image transform
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor()
+])
+
+
+def predict(image_path, text):
+    image = transform(Image.open(image_path).convert("RGB")).unsqueeze(0).to(device)
+    tokens = tokenizer(text, return_tensors="pt", padding=True, truncation=True).to(device)
+
+    with torch.no_grad():
+        img_feat, txt_feat = model(image, tokens["input_ids"], tokens["attention_mask"])
+        similarity = torch.matmul(img_feat, txt_feat.T).squeeze()
+
+    return similarity.item()