my-chatbot / app.py
daksh00's picture
Update app.py
d0e9712 verified
Raw
History Blame Contribute Delete
1.2 kB
import torch
import gradio as gr
import spaces
from transformers import pipeline
pipe = pipeline(
"text-generation",
model="Bouquets/StrikeGPT-R1-Zero-8B",
dtype=torch.bfloat16,
device_map="auto",
)
SYSTEM_PROMPT = (
"You are a helpful AI assistant. "
"Always provide complete, detailed and well-structured answers."
)
@spaces.GPU
def chat(message, history):
# Gradio history is ignored here for simplicity and compatibility.
messages = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": message},
]
prompt = pipe.tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True,
)
output = pipe(
prompt,
max_new_tokens=2048,
temperature=0.6,
top_p=0.9,
do_sample=True,
repetition_penalty=1.1,
return_full_text=False,
pad_token_id=pipe.tokenizer.eos_token_id,
)
return output[0]["generated_text"].strip()
demo = gr.ChatInterface(
fn=chat,
title="🤖 StrikeGPT-R1 Assistant",
description="Powered by Bouquets/StrikeGPT-R1-Zero-8B",
)
if __name__ == "__main__":
demo.launch()