root
commited on
Commit
•
85bfdae
1
Parent(s):
32a066b
add custom pipeline
Browse files- __pycache__/pipeline.cpython-310.pyc +0 -0
- pipeline.py +33 -0
- requirements.txt +1 -0
__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
|