root commited on
Commit
59bec0a
1 Parent(s): 09c3d85

add custom pipeline

Browse files
Files changed (2) hide show
  1. distilbert-base-uncased-emotion +1 -0
  2. handler.py +44 -0
distilbert-base-uncased-emotion ADDED
@@ -0,0 +1 @@
 
 
1
+ Subproject commit d12ff2a4b521b7bfd526aa7055665815c67e113b
handler.py ADDED
@@ -0,0 +1,44 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Dict, Any
2
+ from transformers import pipeline
3
+ import holidays
4
+ import PIL.Image
5
+ import io
6
+ import pytesseract
7
+
8
+ class PreTrainedPipeline():
9
+ def __init__(self, model_path="PrimWong/layout_qa_hparam_tuning"):
10
+ # Initializing the document-question-answering pipeline with the specified model
11
+ self.pipeline = pipeline("document-question-answering", model=model_path)
12
+ self.holidays = holidays.US()
13
+
14
+ def __call__(self, data: Dict[str, Any]) -> str:
15
+ """
16
+ Process input data for document question answering with optional holiday checking.
17
+
18
+ Args:
19
+ data (Dict[str, Any]): Input data containing an 'inputs' field with 'image' and 'question',
20
+ and optionally a 'date' field.
21
+
22
+ Returns:
23
+ str: The answer to the question or a holiday message if applicable.
24
+ """
25
+ inputs = data.get('inputs', {})
26
+ date = data.get("date")
27
+
28
+ # Check if date is provided and if it's a holiday
29
+ if date and date in self.holidays:
30
+ return "Today is a holiday!"
31
+
32
+ # Process the image and question for document question answering
33
+ image_path = inputs.get("image")
34
+ question = inputs.get("question")
35
+
36
+ # Load and process an image
37
+ image = PIL.Image.open(image_path)
38
+ image_text = pytesseract.image_to_string(image) # Use OCR to extract text
39
+
40
+ # Run prediction (Note: this now uses the extracted text, not the image directly)
41
+ prediction = self.pipeline(question=question, context=image_text)
42
+ return prediction["answer"] # Adjust based on actual output format of the model
43
+
44
+ # Note: This script assumes the use of pytesseract for OCR to process images. Ensure pytesseract is configured properly.