import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import accelerate

# Load the model and tokenizer
model_name = "Telugu-LLM-Labs/Indic-gemma-2b-finetuned-sft-Navarasa-2.0"
accelerator = accelerate.Accelerator()
model = AutoModelForCausalLM.from_pretrained(model_name, load_in_4bit=False, device_map="auto", offload_folder="/tmp")
model = accelerator.prepare(model)
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")
    with accelerator.distribute_inputs_to_prepared(model.device_map, inputs):
        outputs = model.generate(**inputs, max_new_tokens=300, use_cache=True)
    response = tokenizer.batch_decode(accelerator.gather(outputs))[0]
    return response.split("### Response:")[-1].strip()

# Gradio interface
with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(label="Input Text")
            instruction = gr.Textbox(label="Instruction", value="Identify and summarize the core intent in the same language:")
        output_text = gr.Textbox(label="Response")
    input_btn = gr.Button("Submit")
    input_btn.click(get_model_response, inputs=[input_text, instruction], outputs=output_text)

demo.launch()