Spaces:
Sleeping
Sleeping
File size: 3,093 Bytes
5651a15 5387ea1 5651a15 5387ea1 5651a15 5387ea1 5651a15 5387ea1 |
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 |
import gradio as gr
import os
import torch
from unsloth import FastLanguageModel
from huggingface_hub import spaces
# Get Hugging Face token from environment variables
HF_TOKEN = os.environ.get('HF_TOKEN')
# Check if we're running in a Hugging Face Space with GPU constraints
IS_SPACES_ZERO = os.environ.get("SPACES_ZERO_GPU", "0") == "1"
IS_SPACE = os.environ.get("SPACE_ID", None) is not None
# Determine device (use GPU if available)
device = "cuda" if torch.cuda.is_available() else "cpu"
LOW_MEMORY = os.getenv("LOW_MEMORY", "0") == "1"
print(f"Using device: {device}")
print(f"Low memory mode: {LOW_MEMORY}")
# Model configuration
max_seq_length = 2048 # Max sequence length for RoPE scaling
dtype = torch.float16 if device == "cuda" else torch.float32
load_in_4bit = True # Enable 4-bit quantization if memory is limited
# Load model and tokenizer with device mapping
model_name = "nafisneehal/chandler_bot"
model, tokenizer = FastLanguageModel.from_pretrained(
model_name=model_name,
max_seq_length=max_seq_length,
dtype=dtype,
load_in_4bit=load_in_4bit,
device_map="auto" if device == "cuda" else None # Automatic GPU mapping
)
FastLanguageModel.for_inference(model) # Optimize model for faster inference
# Define prompt structure (update if necessary for your model)
alpaca_prompt = "{instruction} {input} {output}"
instruction_text = "Learn how to talk like Chandler - a popular character from FRIENDS TV Show. Input is someone saying something, Output is what Chandler saying in response."
@spaces.GPU # Use GPU provided by Hugging Face Spaces if available
def generate_response(user_input, chat_history):
instruction = user_input # Treats user input as instruction
input_text = "" # Any additional input if needed; empty otherwise
# Prepare inputs for model inference on the correct device
inputs = tokenizer(
[alpaca_prompt.format(instruction, input_text, "")],
return_tensors="pt"
).to(device) # Ensure tensors are on the correct device
# Generate response on GPU or CPU as appropriate
with torch.no_grad():
outputs = model.generate(**inputs, max_new_tokens=100)
# Decode response
bot_reply = tokenizer.decode(outputs[0], skip_special_tokens=True)
# Update chat history with user and bot interactions
chat_history.append(("User", user_input))
chat_history.append(("Bot", bot_reply))
return chat_history, "" # Returns updated chat history and clears input
# Set up Gradio interface
with gr.Blocks() as demo:
gr.Markdown("# Llama-Based Chatbot on GPU")
chat_history = gr.Chatbot(label="Chat History")
user_input = gr.Textbox(
placeholder="Type your message here...", label="Your Message")
# Connect submit actions to generate response function
user_input.submit(generate_response, [user_input, chat_history], [
chat_history, user_input])
submit_btn = gr.Button("Send")
submit_btn.click(generate_response, [user_input, chat_history], [
chat_history, user_input])
demo.launch()
|