import gradio as gr
from transformers import pipeline

pipe = pipeline("text-generation", model="keyfan/grok-1-hf", trust_remote_code=True)

def chatbot_response(message, history, max_tokens, temperature):
  history.append(message)
  inp = " ".join(history)
  output = pipe(inp, max_length=max_tokens, temperature=temperature)[0]['generated_text']
  history.append(output)
  return output

iface = gr.ChatInterface(
    fn=chatbot_response, 
    chatbot="Grok-1", 
    history=[], 
    title="Grok-1",
    inputs=[
        gr.Textbox(lines=1, placeholder="Enter your message"),
        gr.Slider(minimum=1, maximum=8192, step=1, value=500, label="Max Tokens"),
        gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.01, label="Temperature")
    ]
)
iface.launch()