Update app.py

app.py

@@ -1,18 +1,45 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import spaces
 
-model_id = "mradermacher/nsfw_chat_0124-GGUF"
+MODEL_ID = "bmiller22000/xyntrai-mistral-2.5-7b-chat-nsfw"
 
-tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto", trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.float16,
+    device_map="auto",
+    trust_remote_code=True
+)
+
+@spaces.GPU(duration=60)  # cấp GPU tạm cho 1 phút
+def chat_with_model(prompt):
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID,
+        torch_dtype=torch.float16,
+        device_map="auto",
+        trust_remote_code=True
+    )
 
-def chat(prompt):
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    outputs = model.generate(**inputs, max_new_tokens=200)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    output = model.generate(
+        **inputs,
+        max_new_tokens=256,
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.9
+    )
+    text = tokenizer.decode(output[0], skip_special_tokens=True)
+    return text
 
-demo = gr.Interface(fn=chat, inputs="text", outputs="text", title="NSFW Chat 0124 (local)")
+demo = gr.Interface(
+    fn=chat_with_model,
+    inputs=gr.Textbox(label="Prompt"),
+    outputs="text",
+    title="Xyntrai-Mistral 2.5 7B Chat (ZeroGPU)"
+)
 
 if __name__ == "__main__":
     demo.launch()