import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# Load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2", device="cpu", trust_remote_code=True)

# Streamlit app
st.title("Text Generation with Transformers")

# Input prompt
prompt = st.text_input("Enter your prompt:")

# Generate button
if st.button("Generate"):
    with torch.no_grad():
        # Tokenize and generate output
        token_ids = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")
        output_ids = model.generate(
            token_ids.to(model.device),
            max_new_tokens=512,
            do_sample=True,
            temperature=0.1
        )
    
    # Decode and display the generated text
    generated_text = tokenizer.decode(output_ids[0][token_ids.size(1):])
    st.text("Generated Text:")
    st.text(generated_text)