Spaces:
Paused
Paused
File size: 2,270 Bytes
d5c810f b21da5b d5c810f b21da5b |
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 |
import streamlit as st
from transformers import AutoTokenizer
from auto_gptq import AutoGPTQForCausalLM
import torch
import subprocess
# Function to get memory info
def get_gpu_memory():
try:
result = subprocess.check_output(["nvidia-smi", "--query-gpu=memory.free,memory.total", "--format=csv,nounits,noheader"], text=True)
memory_info = [x.split(',') for x in result.strip().split('\n')]
memory_info = [{"free": int(x[0].strip()), "total": int(x[1].strip())} for x in memory_info]
except FileNotFoundError:
memory_info = [{"free": "N/A", "total": "N/A"}]
return memory_info
# Display GPU memory information
gpu_memory = get_gpu_memory()
st.write(f"GPU Memory Info: {gpu_memory}")
# Define pretrained model directory
pretrained_model_dir = "FPHam/Jackson_The_Formalizer_V2_13b_GPTQ"
# Check if CUDA is available and get the device
device = "cuda:0" if torch.cuda.is_available() else "cpu"
# Before allocating or loading the model, clear up memory if CUDA is available
if device == "cuda:0":
torch.cuda.empty_cache()
# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(pretrained_model_dir, use_fast=True)
# Attempt to load the model, catch any OOM errors
try:
model = AutoGPTQForCausalLM.from_quantized(
pretrained_model_dir,
model_basename="Jackson2-4bit-128g-GPTQ",
use_safetensors=True,
device=device,
max_memory={0: "10GIB"}
)
except RuntimeError as e:
if 'CUDA out of memory' in str(e):
st.error("CUDA out of memory. Try reducing the model size or input length.")
st.stop()
else:
raise e
# User input for the model
user_input = st.text_input("Input a phrase")
# Generate button
if st.button("Generate the prompt"):
try:
prompt_template = f'USER: {user_input}\nASSISTANT:'
inputs = tokenizer(prompt_template, return_tensors='pt', max_length=512, truncation=True, padding='max_length')
output = model.generate(**inputs)
st.markdown(f"**Generated Text:**\n{tokenizer.decode(output[0])}")
except RuntimeError as e:
if 'CUDA out of memory' in str(e):
st.error("CUDA out of memory during generation. Try reducing the input length.")
else:
raise e |