import gradio as gr
from transformers import pipeline
import os

access_token = os.getenv("Token")
pipe = pipeline("text-generation", model="meta-llama/Llama-3.2-1B-Instruct", token=access_token)

def predict(system="you are a helpful assistant", query):
    messages = []
    messages.append({"role": "system","content": system})
    messages.append({"role": "system","content": query})
    return pipe(messages,max_new_tokens=256)[0]["generated_text"][-1]

gradio_app = gr.Interface(
    predict,
    inputs=["text","text"], outputs="text",
    title="Gemma 2B",
)

if __name__ == "__main__":
    gradio_app.launch()