import gradio as gr
import torch
import os
import torchvision
from transformers import pipeline
auth_token = os.environ.get("HUGGING_FACE_HUB_TOKEN")

# Initialize the pipeline
pipe = pipeline(
    "text-generation",
    model="Jyotiyadav/mistral_7B_NER",
    torch_dtype=torch.bfloat16,
    device_map="auto",
    use_auth_token=auth_token)

# 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
iface = gr.Interface(fn=generate_output, 
                     inputs=["text"],
                     outputs="text", 
                     title="Information Extraction with Mistral-7B",
                     debug=True,
                     description="Generate a Information Extraction with OpenLLM.")
iface.launch()