Hugging Face's logo

Spaces:
hugging-science
/
Mira-API-free
Running on Zero

App Files Community
Mira-API-free
File size: 5,856 Bytes
3bb28fd
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
import gradio as gr
from gradio import Server
import spaces
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# ─────────────────────────────────────────────
# 1. MODEL SETUP
# ─────────────────────────────────────────────
MODEL_ID = "Smilyai-labs/Mira-1-large"

print(f"Loading tokenizer: {MODEL_ID}")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)

print(f"Loading model: {MODEL_ID}")
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float16,
    device_map="auto",       # ZeroGPU manages CUDA device
    trust_remote_code=True,  # Needed for Qwen-based custom archs
)
model.eval()

# ─────────────────────────────────────────────
# 2. INFERENCE FUNCTION (ZeroGPU decorated)
# ─────────────────────────────────────────────
@spaces.GPU(duration=120)
def generate(
    prompt: str,
    system_prompt: str = "You are a helpful assistant.",
    max_new_tokens: int = 512,
    temperature: float = 0.7,
    top_p: float = 0.9,
    do_sample: bool = True,
) -> str:
    """
    Generate a text response from Mira-1-Large.

    Args:
        prompt: The user message / prompt to send to the model.
        system_prompt: System-level instruction for the model.
        max_new_tokens: Maximum number of tokens to generate.
        temperature: Sampling temperature (higher = more creative).
        top_p: Nucleus sampling probability mass.
        do_sample: Whether to use sampling (True) or greedy decoding (False).

    Returns:
        The model's text response as a string.
    """
    # Build chat-style messages (Qwen uses apply_chat_template)
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user",   "content": prompt},
    ]

    # Qwen / Mira chat template
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )

    inputs = tokenizer(text, return_tensors="pt").to(model.device)

    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_p=top_p,
            do_sample=do_sample,
            pad_token_id=tokenizer.eos_token_id,
        )

    # Decode only the newly generated tokens
    new_tokens = output_ids[0][inputs["input_ids"].shape[1]:]
    response = tokenizer.decode(new_tokens, skip_special_tokens=True)
    return response


# ─────────────────────────────────────────────
# 3. STREAMING INFERENCE (SSE / token-by-token)
# ─────────────────────────────────────────────
@spaces.GPU(duration=120)
def generate_stream(
    prompt: str,
    system_prompt: str = "You are a helpful assistant.",
    max_new_tokens: int = 512,
    temperature: float = 0.7,
    top_p: float = 0.9,
) -> str:
    """
    Stream a text response token-by-token from Mira-1-Large via SSE.

    Args:
        prompt: The user message / prompt.
        system_prompt: System-level instruction for the model.
        max_new_tokens: Maximum number of tokens to generate.
        temperature: Sampling temperature.
        top_p: Nucleus sampling probability mass.

    Yields:
        Partial response strings, growing with each new token.
    """
    from transformers import TextIteratorStreamer
    from threading import Thread

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user",   "content": prompt},
    ]
    text = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = tokenizer(text, return_tensors="pt").to(model.device)

    streamer = TextIteratorStreamer(
        tokenizer, skip_prompt=True, skip_special_tokens=True
    )

    gen_kwargs = dict(
        **inputs,
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        top_p=top_p,
        do_sample=True,
        streamer=streamer,
        pad_token_id=tokenizer.eos_token_id,
    )

    thread = Thread(target=model.generate, kwargs=gen_kwargs)
    thread.start()

    partial = ""
    for new_text in streamer:
        partial += new_text
        yield partial


# ─────────────────────────────────────────────
# 4. gr.Server — REST API + OPTIONAL SWAGGER UI
# ─────────────────────────────────────────────
app = Server(
    title="Mira-1-Large API",
    summary="ZeroGPU-backed REST API for Smilyai-labs/Mira-1-large (Qwen arch)",
    version="1.0.0",
)

# Register as Gradio API endpoints (queued, SSE-streaming capable)
app.api(generate,        name="generate")        # POST /gradio_api/call/generate
app.api(generate_stream, name="generate_stream") # POST /gradio_api/call/generate_stream

# Optional: plain FastAPI GET health-check route
@app.get("/health")
def health():
    return {"status": "ok", "model": MODEL_ID}

# ─────────────────────────────────────────────
# 5. LAUNCH
# ─────────────────────────────────────────────
app.launch()