Spaces:
Running on Zero
Running on Zero
File size: 1,951 Bytes
af10acd 81916d7 1ef2cca 3ae1e92 1ef2cca ba0f91a 3ae1e92 ba0f91a 3ae1e92 ba0f91a 3ae1e92 1ef2cca ba0f91a 3ae1e92 67cf64a 7077ef0 af10acd 7077ef0 ba0f91a 7077ef0 84a3ee0 3ae1e92 6847aa2 1ef2cca 6847aa2 3ae1e92 84a3ee0 ba0f91a 46b7aa7 1ef2cca 84a3ee0 6847aa2 ab5b010 1ef2cca ba0f91a 6847aa2 | 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 | import gradio as gr
import spaces
import torch
from fastapi import FastAPI
from transformers import AutoTokenizer, AutoModelForCausalLM
MODEL_NAME = "LiquidAI/LFM2.5-2.6B"
# -------------------------
# LOAD TOKENIZER + MODEL
# -------------------------
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
MODEL_NAME,
dtype=torch.float16
).to("cuda")
model.eval()
# -------------------------
# CHAT FUNCTION
# -------------------------
@spaces.GPU
def model_chat(message, history):
messages = []
for user_msg, bot_msg in history:
messages.append({"role": "user", "content": user_msg})
messages.append({"role": "assistant", "content": bot_msg})
messages.append({"role": "user", "content": message})
inputs = tokenizer.apply_chat_template(
messages,
add_generation_prompt=True,
tokenize=True,
return_dict=True,
return_tensors="pt",
).to(model.device)
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=256,
do_sample=True,
temperature=0.8,
top_p=0.95,
)
generated = outputs[0][inputs["input_ids"].shape[-1]:]
response = tokenizer.decode(generated, skip_special_tokens=True)
return response
# -------------------------
# FASTAPI ENDPOINT (Spaces auto-serves this)
# -------------------------
app = FastAPI()
@app.post("/hf-chat")
def hf_chat(payload: dict):
message = payload["message"]
history = payload.get("history", [])
return {"response": model_chat(message, history)}
# -------------------------
# GRADIO UI (served on main port)
# -------------------------
demo = gr.ChatInterface(
fn=model_chat,
title="LiquidAI/LFM2.5-2.6B Chat Demo",
description="Chat with the LiquidAI/LFM2.5-2.6B model."
)
def main():
demo.launch()
if __name__ == "__main__":
main() |