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

# Function to generate output based on input
def generate_output(input_text, max_new_tokens, temperature, top_k, top_p, model):
    # Initialize the pipeline
    pipe = pipeline(
        "text-generation",
        model=model,
        torch_dtype=torch.bfloat16,
        device_map="auto"
    )
    
    # 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=max_new_tokens,
        temperature=temperature,
        top_k=top_k,
        top_p=top_p,
        num_return_sequences=1,
    )

    # Return the generated text
    return result[0]['generated_text']


examples = [
'''    
COTIZACION FLETE MARITIMO OC 4500325343 Buongiorno ;
Se mi potete quotare: ;
NOT STACKABLE ;
Delivery terms : FCA Anzano del Parco (Como) Italy ;
Pick up address : ; SOREMA Div. of PREVIERO N.S.R.L. ; Via per Cavolto 17 - 22040 Anzano del Parco-CO, Italy ;
Warehouse opening : From 8.30 to 12 h. // from 13,30 to 16,30 h. ;
PO NO. ; DESCRIPTION ; SIZES cm ; WEIGHT ;
4500325343 ; No. 1 CASE ; 85 x 75 x 30 h ; Kg 315 ;
4500325343 ; No. 1 CASE ; 85 x 75 x 30 h ; Kg 315 ;
4500325343 ; No. 1 CASE ; 85 x 75 x 30 h ; Kg 315 ;
4500325343 ; No. 1 CASE ; 85 x 75 x 30 h ; Kg 315 ;
4500325343 ; No. 1 CASE ; 85 x 75 x 30 h ; Kg 315 ;
4500325343 ; No. 1 CASE ; 85 x 75 x 30 h ; Kg 315 ;
4500325343 ; No. 1 CASE ; 85 x 75 x 40 h ; Kg 350 ;
4500325343 ; No. 1 CASE ; 85 x 75 x 40 h ; Kg 350 ;
DESTINO: Veracruz, Mexico ; Gracias y saludos ;
Ivan Briana ; Sales and Operations Manager ;
Tel: +52 55 7155 6559/7098 6337 ; Mobile: +52 (55) 6963 2573 ;
CERMESONI MEXICO S.A.DE C.V. ; Blvd. Adolfo Lopez Mateos no. 202 ;
Col. San Pedro de los Pinos ;
Del. Alvaro Obregon ;
01180 Ciudad de Mexico ;''']



# Create Gradio inputs
inputs = [
    gr.inputs.Textbox(label="Input Text"),
    gr.inputs.Number(label="Max New Tokens", default=32000),
    gr.inputs.Slider(label="Temperature", minimum=0.0, maximum=1.0, default=0.1, step=0.01),
    gr.inputs.Number(label="Top K", default=0),
    gr.inputs.Number(label="Top P", default=0),
    gr.inputs.Textbox(label="Model", default="Jyotiyadav/mistral_7B_NER")
]

# Create a Gradio interface
iface = gr.Interface(
    fn=generate_output, 
    inputs=inputs,
    outputs="text",
    #examples=examples,
    title="Information Extraction with Mistral-7B",
    description="Generate Information Extraction with OpenLLM.",
    debug=True
)

# Launch the interface
iface.launch()