Spaces:
Sleeping
Sleeping
File size: 3,193 Bytes
37197da 5251795 37197da 5251795 37197da 066dad9 5251795 066dad9 5251795 37197da 5251795 37197da 066dad9 5251795 37197da 5251795 |
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 |
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from spaces import GPU
import logging
# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# Load model & tokenizer
MODEL_NAME = "ubiodee/Test_Plutus"
try:
logger.info("Loading tokenizer with use_fast=False...")
tokenizer = AutoTokenizer.from_pretrained(
MODEL_NAME,
use_fast=False, # Use slow tokenizer to avoid fast tokenizer errors
use_safetensors=True,
trust_remote_code=True, # Allow custom tokenizer code
)
logger.info("Tokenizer loaded successfully.")
except Exception as e:
logger.error(f"Tokenizer loading failed: {str(e)}")
raise
try:
logger.info("Loading model with 8-bit quantization...")
model = AutoModelForCausalLM.from_pretrained(
MODEL_NAME,
device_map="auto", # Automatically map to GPU/CPU
load_in_8bit=True, # Use 8-bit quantization to match model
torch_dtype=torch.bfloat16, # Use bfloat16 for efficiency
use_safetensors=True,
low_cpu_mem_usage=True, # Reduce CPU memory during loading
trust_remote_code=True, # Allow custom model code
)
model.eval()
logger.info("Model loaded successfully.")
except Exception as e:
logger.error(f"Model loading failed: {str(e)}")
raise
# Set pad token if not defined
if tokenizer.pad_token_id is None:
tokenizer.pad_token_id = tokenizer.eos_token_id
logger.info("Set pad_token_id to eos_token_id.")
# Move model to GPU if available
if torch.cuda.is_available():
model.to("cuda")
logger.info("Model moved to GPU.")
else:
logger.warning("No GPU available, using CPU.")
# Response function with GPU decorator
@spaces.GPU
def generate_response(prompt, progress=gr.Progress()):
progress(0.1, desc="Tokenizing input...")
try:
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
progress(0.5, desc="Generating response...")
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=200,
temperature=0.7,
top_p=0.9,
do_sample=True,
eos_token_id=tokenizer.eos_token_id,
pad_token_id=tokenizer.pad_token_id,
)
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
# Remove the prompt from the output
if response.startswith(prompt):
response = response[len(prompt):].strip()
progress(1.0, desc="Done!")
return response
except Exception as e:
logger.error(f"Inference failed: {str(e)}")
return f"Error during generation: {str(e)}"
# Gradio UI
demo = gr.Interface(
fn=generate_response,
inputs=gr.Textbox(label="Enter your prompt", lines=4, placeholder="Ask about Plutus..."),
outputs=gr.Textbox(label="Model Response"),
title="Cardano Plutus AI Assistant",
description="Write Plutus smart contracts on Cardano blockchain."
)
# Launch with queueing
demo.queue(max_size=10).launch(enable_queue=True, max_threads=1) |