Update clipGPT.py

clipGPT.py

@@ -70,11 +70,11 @@ class ClipGPT2Model(nn.Module):
 def generate_beam(
     model,
     tokenizer,
+    entry_length,
+    temperature,
+    embed=None,
     beam_size: int = 10,
     prompt=None,
-    embed=None,
-    entry_length=76,
-    temperature=0.9,
     stop_token: str = ".",
 ):
 
@@ -144,7 +144,7 @@ def generate_beam(
 
 
 
-def generate_caption_clipgpt(img):
+def generate_caption_clipgpt(img, entry_length, temperature):
 
     prefix_length = 10
     model = ClipGPT2Model(prefix_length)
@@ -164,7 +164,7 @@ def generate_caption_clipgpt(img):
     with torch.no_grad():
         prefix = clip_model.encode_image(image).to(device, dtype=torch.float32)
         prefix_embed = model.clip_project(prefix).reshape(1, prefix_length, -1)
-        beam_caption = generate_beam(model, tokenizer, embed=prefix_embed)[0]
+        beam_caption = generate_beam(model, tokenizer, entry_length, temperature, embed=prefix_embed)[0]
 
     end_time = time.time()
     print("--- Time taken to generate: %s seconds ---" % (end_time - start_time))