Update app.py

app.py

@@ -6,7 +6,7 @@ from transformers import pipeline
 auth_token = os.environ.get("HUGGING_FACE_HUB_TOKEN")
 
 # Function to generate output based on input
-def generate_output(prompt,max_new_tokens, temperature, top_k, top_p, model):
+def generate_output(input_text,max_new_tokens, temperature, top_k, top_p, model):
     # Initialize the pipeline
     pipe = pipeline(
         "text-generation",
@@ -15,6 +15,39 @@ def generate_output(prompt,max_new_tokens, temperature, top_k, top_p, model):
         torch_dtype=torch.bfloat16,
         device_map="auto"
     )
+
+    prompt = f"""
+    Your task is to extract the information corresponding to the provided labels from the below given email.
+    
+    ### Instruction:
+    * pickup_location: Street address of the origin location of goods.
+    * pickup_cap: Postal code or ZIP code of the pickup location.
+    * pickup_port: Port of pickup, often used in international shipping.
+    * pickup_state: Only Country of pickup location.
+    * delivery_location: Street address of the destination location of goods.
+    * delivery_cap: Postal code or ZIP code of delivery location.
+    * delivery_port: Port of delivery, similar to pickup port.
+    * delivery_state: State or region of delivery location.
+    * total_quantity: Overall quantity of shipped items (e.g., pieces, boxes). Calculate the total_quantity by summing the quantity of all packages.
+    * total_weight: Total weight of the shipment (e.g., kg, lbs). Calculate the total_weight by summing the weights of all packages.
+    * total_volume: Total volume of the shipment (e.g., cubic meters, cubic feet). Calculate the total_volume by summing the volumes of all packages.
+    * quantity: Individual Quantity of a specific item being shipped.
+    * package_type: Individual Type of packaging used (e.g., pallets, cartons).
+    * weight: Individual Weight of a specific package.
+    * measures: Individual Dimensions or measurements of a package.
+    * stackable: Indicates whether the shipment is stackable (True or False).
+    * volume: Individual Volume of a specific package.
+    * commodity: Type of goods or commodities being shipped.
+    * company: Name of the email sending company, also the shipping company or carrier.
+    * incoterms: Choose available options: EXW, FCA, FAS, FOB, CFR, CIF, CPT, CIP, DAP, DPU, DDP.
+    
+    For attributes with multiple values, such as measures, volume, weight, package_type, and quantity, provide each value separately in a JSON format.
+    
+    ### Input:
+    ```{input_text}```
+    
+    ### Response:
+    """
     
     # Generate the result
     result = pipe(