import streamlit as st
import transformers
from transformers import AutoTokenizer
import torch
model = "PY007/TinyLlama-1.1B-Chat-v0.1"
tokenizer = transformers.AutoTokenizer.from_pretrained(model)
pipeline = transformers.pipeline("text-generation",model=model,torch_dtype=torch.bfloat16)

prompt = "What is 6534+2343?"
formatted_prompt = (
    f"### Human: {prompt}### Assistant:"
)


sequences = pipeline(
    formatted_prompt,
    do_sample=True,
    top_k=50,
    top_p = 0.7,
    num_return_sequences=1,
    repetition_penalty=1.1,
    max_new_tokens=500,
    torch_dtype=torch.bfloat16
)
for seq in sequences:
    st.write(f"Result: {seq['generated_text']}")