Update app.py

app.py

@@ -0,0 +1,70 @@
+import gradio as gr
+from transformers import pipeline
+
+# Initialize the pipeline
+pipe = pipeline(
+    "text-generation",
+    model="DataIntelligenceTeam/mistral_7B_NER",
+    torch_dtype=torch.bfloat16,
+    device_map="auto"
+)
+
+# Function to generate output based on input
+def generate_output(input_text):
+    # Prompt for extracting information
+    prompt = f'''
+    Your task is to extract the information corresponding to the provided labels from the below given email.
+    
+    ### Labels:
+    * 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 data:
+    {input_text}
+    
+    ### Output:
+    '''
+
+    # Generate the result
+    result = pipe(
+        f"<s>[INST] {prompt} [/INST]",
+        do_sample=True,
+        max_new_tokens=32000,
+        temperature=0.1,
+        top_k=0,
+        top_p=0,
+        num_return_sequences=1,
+    )
+
+    # Return the generated text
+    return result[0]['generated_text']
+
+# Create a Gradio interface
+gr.Interface(
+    fn=generate_output,
+    inputs=gr.inputs.Textbox(lines=30, label="Input Email"),
+    outputs=gr.outputs.Textbox(lines=30, label="Extracted Information"),
+    title="Email Information Extraction",
+    description="Extract information from emails based on provided labels.",
+    theme="compact"
+).launch()