MPRepo / app.py
AlokSinh's picture
Update app.py
319751b
Raw
History Blame Contribute Delete
2.54 kB
import streamlit as st
from transformers import pipeline
import json
import os
from pprint import pprint
import bitsandbytes as bnb
import torch
import torch.nn as nn
import transformers
from datasets import load_dataset
from huggingface_hub import login,notebook_login
from peft import (
LoraConfig,
PeftConfig,
PeftModel,
get_peft_model,
prepare_model_for_kbit_training
)
from transformers import (
AutoConfig,
AutoModelForCausalLM,
AutoTokenizer,
BitsAndBytesConfig
)
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
PEFT_MODEL = "DhyeyBhalani/mistral_trained_1k"
# notebook_login()
login("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")
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16
)
config = PeftConfig.from_pretrained(PEFT_MODEL)
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)
generation_config = model.generation_config
generation_config.max_new_tokens = 200
generation_config.temperature = 0.7
generation_config.top_p = 0.7
generation_config.num_return_sequences = 1
generation_config.pad_token_id = tokenizer.eos_token_id
generation_config.eos_token_id = tokenizer.eos_token_id
print("--------------------------------------------------------------------------------------------------------------------------------------")
print("--------------------------------------------------------------------------------------------------------------------------------------")
print("-------------------------------------------------------------DONE NO ERROR-------------------------------------------------------------")
print("--------------------------------------------------------------------------------------------------------------------------------------")
print("--------------------------------------------------------------------------------------------------------------------------------------")
# --------------------------------------------------------------------------------------------------------------------------------------
pipe = pipeline('sentiment-analysis')
text = st.text_area('enter some test! 5')
if text:
out = pipe(text)
st.json(out)