root commited on
Commit
85bfdae
1 Parent(s): 32a066b

add custom pipeline

Browse files
__pycache__/pipeline.cpython-310.pyc ADDED
Binary file (1.47 kB). View file
 
pipeline.py ADDED
@@ -0,0 +1,33 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Dict, Any
2
+ from transformers import pipeline
3
+ import holidays
4
+ import PIL.Image
5
+ import io
6
+
7
+ class PreTrainedPipeline():
8
+ def __init__(self, model_path="PrimWong/layout_qa_hparam_tuning"):
9
+ # Initialize the document-question-answering pipeline with the specified model
10
+ self.pipeline = pipeline("document-question-answering", model=model_path)
11
+ self.holidays = holidays.US()
12
+
13
+ def __call__(self, data: Dict[str, Any]) -> str:
14
+ """
15
+ Process input data for document question answering with optional holiday checking.
16
+
17
+ Args:
18
+ data (Dict[str, Any]): Input data containing a 'text' field possibly along with 'image',
19
+ and optionally a 'date' field.
20
+
21
+ Returns:
22
+ str: The answer or processed information based on the text, or a holiday message if applicable.
23
+ """
24
+ text = data.get("inputs")
25
+ date = data.get("date")
26
+
27
+ # Check if the date is a holiday
28
+ if date and date in self.holidays:
29
+ return "Today is a holiday!"
30
+
31
+ # Run prediction using only the text input
32
+ prediction = self.pipeline(question=text, image="What information do you need?")
33
+ return prediction["answer"] # Adjust based on actual output format of the model
requirements.txt ADDED
@@ -0,0 +1 @@
 
 
1
+ holidays