Update app.py

app.py

@@ -0,0 +1,108 @@
+import os
+import bitsandbytes as bnb
+import pandas as pd
+import torch
+import torch.nn as nn
+import transformers
+
+
+from peft import (
+    LoraConfig,
+    PeftConfig,
+    get_peft_model,
+    prepare_model_for_kbit_training,
+    PeftModel
+)
+from transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    BitsAndBytesConfig,
+)
+
+import gradio as gr
+
+import warnings
+
+warnings.filterwarnings("ignore")
+device = "cuda:0"
+
+
+MODEL_NAME = 'diegi97/dolly-v2-6.9b-sharded-bf16'
+
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    load_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16,
+)
+
+model =AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    device_map="auto",
+    trust_remote_code=True,
+    quantization_config=bnb_config,
+)
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+tokenizer.pad_token = tokenizer.eos_token
+
+peft_model_id = "AdiOO7/Azure-Classifier-dolly-7B"
+# peft_model_id = "SparkExpedition/Ticket-Classifier-dolly-7B"
+config = PeftConfig.from_pretrained(peft_model_id)
+model = AutoModelForCausalLM.from_pretrained(
+    config.base_model_name_or_path,
+    return_dict=True,
+    quantization_config=bnb_config,
+    device_map="auto",
+    trust_remote_code=True,
+)
+tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
+tokenizer.pad_token = tokenizer.eos_token
+
+model = PeftModel.from_pretrained(model, peft_model_id)
+
+generation_config = model.generation_config
+generation_config.max_new_tokens = 8
+generation_config.num_return_sequences = 1
+generation_config.temperature = 0.3
+generation_config.top_p = 0.7
+generation_config.pad_token_id = tokenizer.eos_token_id
+generation_config.eos_token_id = tokenizer.eos_token_id
+
+instruct = "From which azure service the issue is raised from {Power BI/Azure Data Factory/Azure Analysis Services}"
+
+def generate_response(question: str) -> str:
+
+    prompt = f"""
+    ### <instruction>: {instruct}
+    ### <human>: {question}
+    ### <assistant>:
+    """.strip()
+
+    encoding = tokenizer(prompt, return_tensors="pt").to(device)
+    with torch.inference_mode():
+        outputs = model.generate(
+            input_ids=encoding.input_ids,
+            attention_mask=encoding.attention_mask,
+            generation_config=generation_config,
+        )
+    response = tokenizer.decode(outputs[0],skip_special_tokens=True)
+
+    assistant_start = '<assistant>:'
+    response_start = response.find(assistant_start)
+    return response[response_start + len(assistant_start):].strip()
+
+labels = ['PowerBI', 'Azure Data Factory', 'Azure Analysis Services']
+
+def answer_prompt(prompt):
+  response = generate_response(prompt)
+  for lab in labels:
+    if response.find(lab) != -1:
+      return lab
+
+iface = gr.Interface(fn=answer_prompt,
+                     inputs=gr.Textbox(lines=5, label="Enter Your Issue", css={"font-size":"18px"}),
+                     outputs=gr.Textbox(lines=5, label="Generated Answer", css={"font-size":"16px"}))
+
+iface.launch()