app.py

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(prompt,max_new_tokens, temperature, top_k, top_p, model):
    # Initialize the pipeline
    pipe = pipeline(
        "text-generation",
        model=model,
        tokenizer = model,
        torch_dtype=torch.bfloat16,
        device_map="auto"
    )
    
    # Generate the result
    result = pipe(
        f"{prompt}",
        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 = [['''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:dubaiDear Team,;Please quote SEA option for below inquiry;POD :Jebel Ali .;estimated packing details;3300x950x1350 mm;3800x1700x1350 mm;1900x2000x1200 mm;1900x2000x1200 mm;2600x1300x1200 mm;3900x1100x1350 mm;820x780x950 mm;Total weight kg 5000 approx;Not stackable;Marrone Srl;via Rui, 5 - Loc. Prà dei Risi;I - 33080 Zoppola (PN);Warehouse: 08:00 – 12:00 / 13:00 – 16:30;Warehouse Handling + Free days w.e.f 01-01-2024;AED 250 + vat until 10 CBM - AED 350 + Vat 11 CBM on wards;Free time applicable - 7 days only for General cargo from DO readiness & No free time for HAZ CARGO, post which standard storage tariff will apply.;Important Notice;:Considering the ongoing circumstances, please be informed after booking confirmation the rates and schedules may undergo changes without advance notice.;Please also expect vessel delays and changes in transit time.;CSS will not be liable for any claims that arise due to this.;Deepak Unnikrishnan;Sr Sales Coordinator;Consolidated Shipping Service L.L.C.;Office: +971 4 883 1303 | Ext: 1163| Toll Free: 800277 | Mobile: +971501513697|;Dubai | United Arab Emirates;nvo-import5@cssdubai.com | www.cssgroupsite.com;Our Offices: Dubai | Abu Dhabi | Sharjah | Ras-Al-Khaimah | Bahrain | Oman | Qatar | Saudi Arabia | Kuwait | Iraq | Africa | Turkiye | India | Sri Lanka |;Disclaimer: This email and any attachments are confidential and may also be privileged. If you are not the intended recipient, please delete all copies and notify the sender immediately.;Please read our;disclaimer.###Response:''']]

# Create Gradio inputs with dropdown options for model selection
model_options = ["DataIntelligenceTeam/model-NER-Phi", "DataIntelligenceTeam/model-NER-Phi"]
#"DataIntelligenceTeam/mistral_7B_NER"]

# Create Gradio inputs
inputs = [
    gr.inputs.Textbox(label="Input Text"),
    gr.inputs.Number(label="Max New Tokens", default=2000),
    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.Dropdown(label="Model", choices=model_options, default=model_options[0])
]

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

# Launch the interface
iface.launch()