import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# MODEL_ID can be any HF model repo you want to load
MODEL_ID = "meta-llama/Llama-3.2-1B-Instruct"  # replace with your own on the Hub

# load tokenizer + model
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID)

# chat function
def chat(input_text, history=[]):
    # append user input to history
    history = history or []
    history.append(("You", input_text))

    # encode input
    inputs = tokenizer(
        input_text + tokenizer.eos_token,
        return_tensors="pt"
    )

    # generate a response
    output_ids = model.generate(
        **inputs,
        max_new_tokens=100,
        pad_token_id=tokenizer.eos_token_id
    )

    response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    history.append(("Bot", response))

    # return formatted history
    return "", history

# Gradio UI
with gr.Blocks() as demo:
    chatbot = gr.Chatbot()
    user_input = gr.Textbox(placeholder="Type a message…")
    user_input.submit(chat, [user_input, chatbot], [user_input, chatbot])

demo.launch()