Spaces:
Sleeping
Sleeping
File size: 3,035 Bytes
6bb8442 f4a2c82 6bb8442 f4a2c82 6bb8442 f4a2c82 6bb8442 | 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 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | import asyncio
from typing import AsyncIterator, Optional
import spaces
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
MODEL_ID = "Qwen/Qwen2.5-1.5B-Instruct"
@spaces.GPU(duration=120)
def _generate_sync(
messages: list,
max_tokens: int,
) -> str:
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
MODEL_ID,
torch_dtype=torch.float16,
device_map="cuda",
)
print("\n========== CLOUD LLM PROMPT DEBUG ==========")
print(f"Message count: {len(messages)}")
for i, message in enumerate(messages):
role = message.get("role", "UNKNOWN")
content = message.get("content", "")
print(
f"[MESSAGE {i}] "
f"role={role} "
f"chars={len(content)}"
)
print("============================================\n")
prompt = tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True,
)
inputs = tokenizer(
prompt,
return_tensors="pt",
).to("cuda")
print(
f"[CLOUD LLM] Actual input tokens: "
f"{inputs.input_ids.shape[1]}"
)
with torch.inference_mode():
outputs = model.generate(
**inputs,
max_new_tokens=max_tokens,
temperature=0.6,
top_p=0.9,
do_sample=True,
)
generated = outputs[0][inputs.input_ids.shape[1]:]
response = tokenizer.decode(
generated,
skip_special_tokens=True,
)
del outputs
del inputs
del model
del tokenizer
torch.cuda.empty_cache()
return response
class CloudLLMEngine:
def __init__(self):
self.last_usage = None
async def stream_tokens(
self,
user_text: str,
) -> AsyncIterator[str]:
messages = [
{
"role": "system",
"content": (
"You are Edi, an AI engineering mentor."
),
},
{
"role": "user",
"content": user_text,
},
]
async for token in self.stream_tokens_from_messages(messages):
yield token
async def stream_tokens_from_messages(
self,
messages: list,
session_id: str = "",
max_tokens: Optional[int] = None,
) -> AsyncIterator[str]:
response = await asyncio.to_thread(
_generate_sync,
messages,
max_tokens or 250,
)
# Compatibility streaming.
#
# ZeroGPU generation completes before this adapter receives
# the result. We emit small text chunks afterward so the
# existing AgentController can consume an AsyncIterator.
chunk_size = 12
for index in range(0, len(response), chunk_size):
yield response[index:index + chunk_size]
await asyncio.sleep(0)
async def aclose(self):
return |