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

# Load the model and tokenizer
model_name = "rajrakeshdr/IntelliSoc"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Streamlit app title
st.title("IntelliSoc Text Generation")

# Input prompt
prompt = st.text_area("Enter your prompt:", "Once upon a time")

# Slider for max length
max_length = st.slider("Max length of generated text", 50, 200, 100)

# Generate text on button click
if st.button("Generate Text"):
    # Tokenize input
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, padding=True)
    
    # Generate text
    with torch.no_grad():
        outputs = model.generate(
            inputs.input_ids,
            max_length=max_length,
            num_return_sequences=1,
            no_repeat_ngram_size=2,
            top_k=50,
            top_p=0.95,
            temperature=0.7
        )
    
    # Decode the generated text
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # Display the generated text
    st.write("Generated Text:")
    st.write(generated_text)