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

model_name = "ai4bharat/Airavata"
tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left")
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)

SYSTEM_PROMPT = """<s>[INST] <<SYS>>
नमस्कार! आप अब कृषि विशेषज्ञता बॉट के साथ इंटरैक्ट कर रहे हैं—एक उन्नत AI जो कृषि क्षेत्र में विशेषज्ञता प्रदान करने के लिए डिज़ाइन किया गया है।

कृपया ध्यान दें कि यह बॉट केवल हिंदी में जवाब देगा। इसकी क्षमताएँ शामिल हैं:

1. आधुनिक फसल प्रबंधन तकनीकों में गहरा ज्ञान।
2. कृषि में कीट और रोग नियंत्रण के लिए प्रभावी रणनीतियाँ।
3. मृदा स्वास्थ्य का सुधारने और पुनर्निर्माण के लिए विशेषज्ञता।
4. सतत और प्रेसिजन खेती के अभ्यासों का ज्ञान।
5. सिंचाई और जल प्रबंधन के लिए सर्वोत्तम अभ्यासों के लिए सुझाव।
6. रणनीतिक फसल चक्रण और इंटरक्रॉपिंग विधियों पर मार्गदर्शन।
7. नवीनतम कृषि प्रौद्योगिकियों और नवाचारों की जानकारी।
8. विशेष फसलों, जलवायु, और क्षेत्रों के लिए विशेषज्ञ सलाह।

कृपया पेशेवर रूप से बराबरी बनाए रखें और सुनिश्चित करें कि आपके जवाब सही और मूल्यवान हैं। उपयोगकर्ताओं से आगे की स्पष्टीकरण के लिए पूछने के लिए प्रोत्साहित करें।

आपका प्रमुख लक्ष्य है यह है कि आप कृषि क्षेत्र में उपयुक्त ज्ञान प्रदान करें। आपके ज्ञान का धन्यवाद।
<</SYS>>
"""

device = "cuda" if torch.cuda.is_available() else "cpu"

def create_prompt_with_chat_format(messages, bos="<s>", eos="</s>", add_bos=True, system_prompt="System: "):
    formatted_text = ""
    for message in messages:
        if message["role"] == "system":
            formatted_text += system_prompt + message["content"] + "\n"
        elif message["role"] == "user":
            if isinstance(message["content"], list):
                formatted_text += "\n" + "\n".join(message["content"]) + "\n"
            else:
                formatted_text += "\n" + message["content"] + "\n"
        elif message["role"] == "assistant":
            if isinstance(message["content"], list):
                formatted_text += "\n" + "\n".join(message["content"]).strip() + eos + "\n"
            else:
                formatted_text += "\n" + message["content"].strip() + eos + "\n"
        else:
            raise ValueError(
                "Tulu chat template only supports 'system', 'user', and 'assistant' roles. Invalid role: {}.".format(
                    message["role"]
                )
            )
    formatted_text += "\n"
    formatted_text = bos + formatted_text if add_bos else formatted_text
    return formatted_text


def inference(input_prompts, model, tokenizer, system_prompt="System: "):
    output_texts = []
    model = model.to(device)  # Move the model to the same device as the input data
    for input_prompt in input_prompts:
        formatted_query = create_prompt_with_chat_format([{"role": "user", "content": input_prompt}], add_bos=False, system_prompt=system_prompt)
        encodings = tokenizer(formatted_query, padding=True, return_tensors="pt")
        encodings = {key: value.to(device) for key, value in encodings.items()}  # Move input data to the same device as the model

        with torch.no_grad():
            outputs = model.generate(encodings["input_ids"], do_sample=False, max_length=250)

        output_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        output_texts.append(output_text[len(input_prompt):])
    return output_texts




examples = [
    ["मुझे अपने करियर के बारे में सुझाव दो", "मैं कैसे अध्ययन कर सकता हूँ?"],
    ["कृपया मुझे एक कहानी सुनाएं", "ताजमहल के बारे में कुछ बताएं"],
    ["मेरा नाम क्या है?", "आपका पसंदीदा फिल्म कौन सी है?"],
]


def get_llama_response(message: str, history: list, system_prompt=SYSTEM_PROMPT) -> str:
    formatted_history = [{"role": "user", "content": hist} for hist in history]
    formatted_message = {"role": "user", "content": message}

    formatted_query = create_prompt_with_chat_format(formatted_history + [formatted_message], add_bos=False, system_prompt=system_prompt)
    response = inference([formatted_query], model, tokenizer)

    print("Chatbot:", response[0].strip())
    return response[0].strip()


gr.ChatInterface(fn=get_llama_response).launch()