import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM

# Streamlit App Title
st.title("Tamil Text Generation with LLaMA")

# Load the model and tokenizer
model_name = "abhinand/tamil-llama-7b-base-v0.1"
st.sidebar.write("Loading the model... This may take some time.")
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

st.sidebar.write("Model loaded successfully!")

# Text input from the user
input_text = st.text_area("Enter Tamil text:", "வணக்கம், எப்படி இருக்கின்றீர்கள்?")

# Generate button
if st.button("Generate Text"):
    with st.spinner("Generating response..."):
        # Encode the input text
        inputs = tokenizer(input_text, return_tensors="pt")

        # Generate response
        outputs = model.generate(**inputs, max_length=50)

        # Decode and display the generated text
        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        st.text_area("Generated Response:", generated_text, height=200)