Update handler.py

handler.py

@@ -4,7 +4,7 @@ from obsidian import Vision
 
 class EndpointHandler():
     def __init__(self, path="", vision_model="obsidian3b"):
-        self.pipeline = pipeline("text-classification", model=path)
+        self.pipeline = pipeline("text-generation", model=path)
         self.vision = Vision(vision_model)
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
@@ -29,12 +29,12 @@ class EndpointHandler():
             combined_captions = [inputs, image_caption]
 
             # run text classification on combined captions
-            prediction = self.pipeline(combined_captions, temperature=0.33, num_beams=5, stop=[])
+            prediction = self.pipeline(combined_captions, temperature=0.33, num_beams=5, stop=[], do_sample=True)
 
             return prediction
 
         else:
             # run text classification on plain text input
-            prediction = self.pipeline(inputs, temperature=0.33, num_beams=5, stop=[])
+            prediction = self.pipeline(inputs, temperature=0.33, num_beams=5, stop=[], do_sample=True)
 
             return prediction