import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

model_name = "google/flan-t5-xl"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

gita_context = """
The Bhagavad Gita is a 700-verse Hindu scripture that is part of the Indian epic Mahabharata. It is a dialogue between Prince Arjuna and Lord Krishna, who serves as his charioteer. The Gita's core message includes:
1. The immortality of the soul (Atman)
2. The nature of action (Karma) and duty (Dharma)
3. The importance of devotion (Bhakti)
4. The pursuit of knowledge (Jnana) and wisdom
5. Different types of Yoga: Karma Yoga, Bhakti Yoga, Jnana Yoga, and Raja Yoga
6. The concept of detachment from the fruits of one's actions
7. The divine nature of Krishna as an avatar of Vishnu
Key teachings include performing one's duty without attachment to results, the importance of self-realization, and the path to liberation (Moksha).
"""

def generate_response(question):
    prompt = f"Based on the following context about the Bhagavad Gita, answer the question.\n\nContext: {gita_context}\n\nQuestion: {question}\n\nAnswer:"
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
    outputs = model.generate(input_ids, max_new_tokens=200, do_sample=True, temperature=0.7, top_p=0.95)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

iface = gr.Interface(
    fn=generate_response,
    inputs=gr.Textbox(lines=2, placeholder="Enter your question about the Bhagavad Gita here..."),
    outputs="text"
)

iface.launch()