Spaces:
Sleeping
Sleeping
File size: 1,889 Bytes
b157648 f412d3b b157648 f412d3b b157648 |
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 |
# app.py
import streamlit as st
from unsloth import FastLanguageModel
from transformers import TextStreamer
# To speed up model loading in repeated queries, you can use st.cache_resource (Streamlit 1.18+).
@st.cache_resource
def load_unsloth_model(
model_name="azizsi/model2",
max_seq_length=4096,
dtype="float16",
load_in_4bit=False
):
"""
Loads and prepares the model for inference using FastLanguageModel from Unsloth.
Returns (model, tokenizer).
"""
model, tokenizer = FastLanguageModel.from_pretrained(
model_name=model_name,
max_seq_length=max_seq_length,
dtype=dtype,
load_in_4bit=load_in_4bit
)
# Enable 2x faster inference (per Unsloth docs)
FastLanguageModel.for_inference(model)
return model, tokenizer
def main():
st.title("Unsloth Model Demo")
# Provide a text input area for the user
user_input = st.text_area("Enter your prompt:", "")
# Generate button
if st.button("Generate"):
with st.spinner("Generating response..."):
# Load the model & tokenizer
model, tokenizer = load_unsloth_model()
# Create a TextStreamer to stream tokens or capture final text
streamer = TextStreamer(tokenizer)
# Tokenize user prompt and move to GPU (or the model's device)
inputs = tokenizer(user_input, return_tensors="pt").to(model.device)
# Generate up to 128 new tokens (modify as desired)
outputs = model.generate(**inputs, streamer=streamer, max_new_tokens=128)
# If you want to display the entire response at once:
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
st.markdown("**Response:**")
st.write(generated_text)
if __name__ == "__main__":
main()
|