import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load the model and tokenizer
model_name = "Telugu-LLM-Labs/Indic-gemma-2b-finetuned-sft-Navarasa-2.0"
model = AutoModelForCausalLM.from_pretrained(model_name, load_in_4bit=False, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_name)

def generate_prompt(instruction, user_input):
    """
    Generates a prompt for the model to ensure it responds with the intent in the same language as the input.
    """
    return f"""
### Instruction:
{instruction}

### Input:
{user_input}

### Response:
"""

def get_model_response(user_input, instruction="Identify and summarize the core intent in the same language:"):
    """
    Gets the model's response, ensuring it matches the input language and focuses on extracting a concise intent.
    """
    input_text = generate_prompt(instruction, user_input)
    inputs = tokenizer([input_text], return_tensors="pt")
    outputs = model.generate(**inputs, max_new_tokens=300, use_cache=True)
    response = tokenizer.batch_decode(outputs)[0]
    return response.split("### Response:")[-1].strip()

# Gradio interface
iface = gr.Interface(
    fn=get_model_response,
    inputs=[
        gr.inputs.Textbox(label="Input Text"),
        gr.inputs.Textbox(label="Instruction", default="Identify and summarize the core intent in the same language:"),
    ],
    outputs=gr.outputs.Textbox(label="Response"),
    title="Intent Summarization",
    description="Summarize the core intent of the input text in the same language.",
)

iface.launch()