Spaces:
Runtime error
Runtime error
File size: 3,010 Bytes
ae24dbd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 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() |