Spaces:
Sleeping
Sleeping
File size: 3,159 Bytes
6d0b18e | 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 | import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from transformers.cache_utils import DynamicCache
import streamlit as st
# Add necessary serialization safety
torch.serialization.add_safe_globals([DynamicCache])
torch.serialization.add_safe_globals([set])
# Minimal generate function for token-by-token generation
def generate(model, input_ids, past_key_values, max_new_tokens=50):
"""
This function performs token-by-token text generation using a pre-trained language model.
Purpose: To generate new text based on input tokens, without loading the full context repeatedly
Process: It takes a model, input IDs, and cached key-values, then generates new tokens one by one up to the specified maximum
Performance: Uses the cached key-values for efficiency and returns only the newly generated tokens
"""
device = model.model.embed_tokens.weight.device
origin_len = input_ids.shape[-1]
input_ids = input_ids.to(device)
output_ids = input_ids.clone()
next_token = input_ids
with torch.no_grad():
for _ in range(max_new_tokens):
out = model(
input_ids=next_token,
past_key_values=past_key_values,
use_cache=True
)
logits = out.logits[:, -1, :]
token = torch.argmax(logits, dim=-1, keepdim=True)
output_ids = torch.cat([output_ids, token], dim=-1)
past_key_values = out.past_key_values
next_token = token.to(device)
if model.config.eos_token_id is not None and token.item() == model.config.eos_token_id:
break
return output_ids[:, origin_len:] # Return just the newly generated part
def get_kv_cache(model, tokenizer, prompt):
"""
This function creates a key-value cache for a given prompt.
Purpose: To pre-compute and store the model's internal representations (key-value states) for a prompt
Process: Encodes the prompt, runs it through the model, and captures the resulting cache
Returns: The cache object and the original prompt length for future reference
"""
# Encode prompt
device = model.model.embed_tokens.weight.device
input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
cache = DynamicCache() # it grows as text is generated
# Run the model to populate the KV cache:
with torch.no_grad():
_ = model(
input_ids=input_ids,
past_key_values=cache,
use_cache=True
)
return cache, input_ids.shape[-1]
# Initialize session state for the model, tokenizer and cache
@st.cache_resource
def load_model_and_tokenizer():
model_name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
tokenizer = AutoTokenizer.from_pretrained(
model_name,
trust_remote_code=True
)
model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
device_map="auto",
trust_remote_code=True
)
return model, tokenizer |