Update app.py

app.py

@@ -5,91 +5,70 @@ 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,
-        tokenizer = model,
-        torch_dtype=torch.bfloat16,
-        device_map="auto"
-    )
+import gradio as gr
+from unsloth import FastLanguageModel
 
-    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(
-        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,
-    )
+instruction = """
+From the given email, extract the following key values. The keys are explained below:
+* 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.
+"""
 
-    # Return the generated text
-    return result[0]['generated_text']
 
+# Define the function for generating output based on input
+def generate_output(input_text):
+    # Prompt for the instruction
+    
+    output = ""
+    # Initialize the FastLanguageModel
+    model, tokenizer = FastLanguageModel.from_pretrained(
+        model_name = "sxandie/llama_3_8b_4bitQ",
+        max_seq_length = 2048,
+        dtype = None,
+        load_in_4bit = True,
+    )
+    FastLanguageModel.for_inference(model) # Enable native 2x faster inference
 
-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;[email protected] | 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:''']]
+    alpaca_prompt = f"""
+    Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
 
-# Create Gradio inputs with dropdown options for model selection
-model_options = ["DataIntelligenceTeam/model-NER-Phi", "DataIntelligenceTeam/mistral_7B_NER"]
-#"DataIntelligenceTeam/mistral_7B_NER"]
+    ### Instruction:
+    {instruction}
 
-# 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])
-]
+    ### Input:
+    {input_text}
 
-# 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
-)
+    ### Response:
+    """
+    # Tokenize the input text
+    inputs = tokenizer([alpaca_prompt], return_tensors="pt").to("cuda")
+    # Generate outputs
+    outputs = model.generate(**inputs, max_new_tokens=2048, use_cache=True)
+    output = tokenizer.batch_decode(outputs)
+    return output
 
-# Launch the interface
+# Create Gradio interface
+iface = gr.Interface(fn=generate_output, 
+                     inputs="text", 
+                     outputs="text", 
+                     title="Email Information Extraction", 
+                     description="Extract key information from the provided email.")
 iface.launch()