app.py

import re
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# =====================================================
# MODEL PATH
# =====================================================
MODEL_PATH = r"D:\agriculture_chatbot\Agriculture_Ai_Assistant\model_k_1"

# =====================================================
# LOAD TOKENIZER
# =====================================================
print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(
    MODEL_PATH,
    trust_remote_code=True,
)

# =====================================================
# LOAD MODEL
# =====================================================
print("Loading model...")

model = AutoModelForCausalLM.from_pretrained(
    MODEL_PATH,
    torch_dtype=torch.float16,
    device_map="auto",
    trust_remote_code=True,
)

model.eval()

print(f"Model loaded on: {model.device}")

# Optional (PyTorch 2.x)
try:
    model = torch.compile(model)
    print("Model compiled successfully.")
except Exception:
    pass


# =====================================================
# LANGUAGE DETECTION
# =====================================================
def is_gujarati(text):
    return bool(re.search(r'[\u0A80-\u0AFF]', text))


# =====================================================
# LONG ANSWER DETECTION
# =====================================================
def needs_long_answer(question):

    q = question.lower()

    keywords = [
        "explain",
        "describe",
        "detail",
        "detailed",
        "why",
        "how",
        "advantages",
        "disadvantages",
        "compare",
        "difference",
        "process",
        "steps",
        "cultivation",
        "management",
        "control",
        "treatment",

        # Gujarati
        "સમજાવો",
        "વિગત",
        "વિસ્તાર",
        "શા માટે",
        "કેવી રીતે",
        "ફાયદા",
        "ગેરફાયદા",
        "તફાવત",
        "પ્રક્રિયા",
        "પદ્ધતિ"
    ]

    return any(k in q for k in keywords)


# =====================================================
# SYSTEM PROMPTS
# =====================================================
ENGLISH_SYSTEM = """
You are Agriculture AI Assistant.

You are an expert in:

- Agriculture
- Crops
- Fertilizers
- Soil Science
- Irrigation
- Plant Diseases
- Insects
- Weather
- Farming

Rules:

1. Reply ONLY in English.
2. Never use Gujarati.
3. Never invent facts.
4. If you don't know, clearly say:
   "I don't have enough information."
5. Give accurate agricultural information.
6. Give short answers unless the user asks for explanation.
7. If the user asks to explain, provide detailed headings, bullet points, examples and conclusion.
"""


GUJARATI_SYSTEM = """
તમે કૃષિ AI સહાયક છો.

નિયમો:

1. માત્ર ગુજરાતી ભાષામાં જવાબ આપો.
2. અંગ્રેજી શબ્દોનો ઉપયોગ ન કરો.
3. ખોટી માહિતી આપશો નહીં.
4. માહિતી ન હોય તો સ્પષ્ટ કહો કે માહિતી ઉપલબ્ધ નથી.
5. ટૂંકો જવાબ આપો.
6. જો વપરાશકર્તા વિગત માંગે તો સંપૂર્ણ વિગતવાર જવાબ આપો.
"""


# =====================================================
# ASK FUNCTION
# =====================================================
def ask(question):

    if is_gujarati(question):
        system_prompt = GUJARATI_SYSTEM
    else:
        system_prompt = ENGLISH_SYSTEM

    long_answer = needs_long_answer(question)

    if long_answer:
        max_tokens = 1024

        question += """

Please provide:

- Introduction
- Detailed explanation
- Important points
- Practical recommendations
- Conclusion
"""

    else:
        max_tokens = 256

    messages = [
        {
            "role": "system",
            "content": system_prompt
        },
        {
            "role": "user",
            "content": question
        }
    ]

    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )

    inputs = tokenizer(
        prompt,
        return_tensors="pt"
    ).to(model.device)

    with torch.inference_mode():

        outputs = model.generate(
            **inputs,

            max_new_tokens=max_tokens,

            do_sample=False,

            repetition_penalty=1.08,

            no_repeat_ngram_size=3,

            use_cache=True,

            early_stopping=True,

            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )

    response = tokenizer.decode(
        outputs[0][inputs.input_ids.shape[1]:],
        skip_special_tokens=True
    )

    return response.strip()


# =====================================================
# MAIN
# =====================================================
if __name__ == "__main__":

    print("=" * 65)
    print("          Agriculture AI Assistant")
    print("          English + Gujarati")
    print("          Type 'quit' to exit")
    print("=" * 65)

    while True:

        question = input("\nYou : ").strip()

        if not question:
            continue

        if question.lower() in ["quit", "exit", "q"]:
            print("\nGoodbye!")
            break

        print("\nAssistant:\n")

        answer = ask(question)

        print(answer)