import streamlit as st
from transformers import pipeline
import torch
import os
from huggingface_hub import login


login(token = os.getenv('HF_TOKEN'))
pipe = pipeline("text-generation", model="google/gemma-3-1b-it", torch_dtype=torch.bfloat16)

txt = st.text_area("Test this")

if (txt):
    messages = [
        [
            {
                "role": "system",
                "content": [{"type": "text", "text": "You are a helpful assistant."},]
            },
            {
                "role": "user",
                "content": [{"type": "text", "text": "Write a poem on Hugging Face, the company"},]
            },
        ],
    ]
    st.write(pipe(messages, max_new_tokens=50))