Create inference.py

inference.py

@@ -0,0 +1,31 @@
+from transformers import pipeline, Text2TextGenerationPipeline
+
+model_name = "test-BB2"
+model = pipeline(
+    "text2text-generation",
+    model=model_name,
+    tokenizer=model_name,
+)
+
+class MyText2TextGenerationPipeline(Text2TextGenerationPipeline):
+    def __init__(self, *args, default_text=None, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.default_text = default_text
+
+    def __call__(self, *args, **kwargs):
+        if "prompt" not in kwargs and self.default_text is not None:
+            kwargs["prompt"] = self.default_text
+        return super().__call__(*args, **kwargs)
+
+def generate_text(input_text, default_text="Enter your input text here"):
+    generator = MyText2TextGenerationPipeline(
+        model=model,
+        tokenizer=model.tokenizer,
+        default_text=default_text
+    )
+    return generator(input_text, max_length=100)
+
+# Example usage:
+input_text = "The quick brown fox jumps over the lazy dog."
+generated_text = generate_text(input_text, default_text="Enter a new input text here")
+print(generated_text)