|
|
import os |
|
|
import gradio as gr |
|
|
import spaces |
|
|
import torch |
|
|
import logging |
|
|
import time |
|
|
|
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM |
|
|
from transformers.utils import logging as hf_logging |
|
|
|
|
|
logging.basicConfig( |
|
|
filename="/tmp/app.log", |
|
|
level=logging.DEBUG, |
|
|
format="%(asctime)s %(levelname)s: %(message)s" |
|
|
) |
|
|
|
|
|
logging.info("Starting app.py logging") |
|
|
hf_logging.set_verbosity_debug() |
|
|
hf_logging.set_verbosity_info() |
|
|
hf_logging.enable_default_handler() |
|
|
hf_logging.enable_explicit_format() |
|
|
hf_logging.add_handler(logging.FileHandler("/tmp/transformers.log")) |
|
|
|
|
|
|
|
|
model_id = "futurehouse/ether0" |
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(model_id) |
|
|
model = AutoModelForCausalLM.from_pretrained( |
|
|
model_id, |
|
|
device_map="auto", |
|
|
torch_dtype=torch.float16 |
|
|
) |
|
|
|
|
|
@spaces.GPU |
|
|
def chat_fn(prompt, max_tokens=512): |
|
|
t0 = time.time() |
|
|
max_tokens = min(int(max_tokens), 32_000) |
|
|
|
|
|
try: |
|
|
messages = [{"role": "user", "content": prompt}] |
|
|
chat_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) |
|
|
|
|
|
inputs = tokenizer(chat_prompt, return_tensors="pt").to(model.device) |
|
|
t1 = time.time() |
|
|
logging.info(f"🧠 Tokenization complete in {t1 - t0:.2f}s") |
|
|
|
|
|
outputs = model.generate( |
|
|
**inputs, |
|
|
max_new_tokens=max_tokens, |
|
|
do_sample=True, |
|
|
temperature=0.1, |
|
|
pad_token_id=tokenizer.eos_token_id |
|
|
) |
|
|
t2 = time.time() |
|
|
logging.info(f"⚡️ Generation complete in {t2 - t1:.2f}s (max_tokens={max_tokens})") |
|
|
|
|
|
generated_text = tokenizer.decode( |
|
|
outputs[0][inputs['input_ids'].shape[1]:], |
|
|
skip_special_tokens=True |
|
|
) |
|
|
t3 = time.time() |
|
|
logging.info(f"🔓 Decoding complete in {t3 - t2:.2f}s (output length: {len(generated_text)})") |
|
|
|
|
|
return generated_text |
|
|
|
|
|
except Exception: |
|
|
logging.exception("❌ Exception during generation") |
|
|
return "⚠️ Generation failed" |
|
|
|
|
|
gr.Interface( |
|
|
fn=chat_fn, |
|
|
inputs=[ |
|
|
gr.Textbox(label="prompt"), |
|
|
gr.Number(label="max_tokens", value=512, precision=0) |
|
|
], |
|
|
outputs="text", |
|
|
title="Ether0" |
|
|
).launch(ssr_mode=False) |
