# app.py
import streamlit as st
from transformers import AutoModelForCausalLM, AutoTokenizer

def generate_kannada_text(prompt):
    model_name = "Tensoic/Kan-LLaMA-7B-SFT-v0.1"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name)

    input_ids = tokenizer.encode(prompt, return_tensors="pt")

    output = model.generate(
        input_ids,
        max_length=150,
        num_beams=5,
        no_repeat_ngram_size=2,
        top_k=50,
        top_p=0.95,
        length_penalty=0.8
    )

    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
    return generated_text

def main():
    st.title("Kannada Text Generation App")

    # User input prompt
    prompt = st.text_area("Enter a prompt in Kannada:")

    # Generate Kannada text
    if st.button("Generate Text"):
        generated_text = generate_kannada_text(prompt)
        st.subheader("Generated Kannada Text:")
        st.write(generated_text)

if __name__ == "__main__":
    main()