Create caption_refiner.py

caption_refiner.py

@@ -0,0 +1,18 @@
+from transformers import T5Tokenizer, T5ForConditionalGeneration
+
+# Load the T5 model and tokenizer
+tokenizer = T5Tokenizer.from_pretrained("t5-base")
+model = T5ForConditionalGeneration.from_pretrained("t5-base")
+
+def refine_caption(initial_caption):
+    # Prepare the input for the T5 model
+    input_text = f"refine caption: {initial_caption}"
+    input_ids = tokenizer(input_text, return_tensors="pt").input_ids
+    
+    # Generate the refined caption
+    outputs = model.generate(input_ids, max_new_tokens=100, num_return_sequences=1)
+    
+    # Decode the refined caption
+    refined_caption = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    return refined_caption