Spaces:
Paused
Paused
| import spaces | |
| import torch | |
| import gradio as gr | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| # --- Model loading block --- | |
| # This is the part that's specific to Nile-Chat. If you ever swap models, | |
| # this is the only section you'd need to change (usually just the model_id). | |
| model_id = "MBZUAI-Paris/Nile-Chat-4B" | |
| tokenizer = AutoTokenizer.from_pretrained(model_id) | |
| model = AutoModelForCausalLM.from_pretrained( | |
| model_id, | |
| dtype=torch.bfloat16, | |
| device_map="auto", | |
| ) | |
| # --- Generation function --- | |
| # ZeroGPU only hands out a GPU for the duration of a function marked | |
| # with @spaces.GPU — everything outside this function runs on CPU. | |
| def generate_response(message, history): | |
| # This Gradio version passes history as a list of [user_msg, bot_msg] | |
| # pairs — convert each pair into the role/content format the model expects. | |
| 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}) | |
| prompt = tokenizer.apply_chat_template( | |
| messages, tokenize=False, add_generation_prompt=True | |
| ) | |
| inputs = tokenizer(prompt, return_tensors="pt").to(model.device) | |
| output = model.generate( | |
| **inputs, | |
| max_new_tokens=512, | |
| temperature=0.7, | |
| do_sample=True, | |
| ) | |
| # Slice off the prompt in token space (not string space) — avoids any | |
| # mismatch between how the prompt and full output get decoded. | |
| input_len = inputs["input_ids"].shape[1] | |
| new_tokens = output[0][input_len:] | |
| reply = tokenizer.decode(new_tokens, skip_special_tokens=True) | |
| return reply.strip() | |
| # --- Chat interface block --- | |
| # This part is boilerplate — the same shape for almost any chat model Space. | |
| demo = gr.ChatInterface( | |
| fn=generate_response, | |
| title="Nile-Chat Egyptian Arabic Assistant", | |
| description="Chat in Egyptian Arabic (Arabic script or Arabizi).", | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() |