import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("ai4bharat/Airavata")
model = AutoModelForCausalLM.from_pretrained("ai4bharat/Airavata")

def generate_response(prompt):
    input_ids = tokenizer.encode(prompt, return_tensors="pt", max_length=50)
    output_ids = model.generate(input_ids, max_length=100, num_beams=5, no_repeat_ngram_size=2)
    response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    return response

iface = gr.Interface(
    fn=generate_response,
    inputs="text",
    outputs="text",
    live=True,
    title="Airavata LLMs Chatbot",
    description="Ask me anything, and I'll generate a response!",
    theme="light",
)

iface.launch()












# 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()