import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread

# Model setup
model_id = "Qwen/Qwen2.5-0.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)

def chat(message, history):
    # Aapka Original System Prompt
    system_prompt = (
        "You are 'Jyoti Agent', the official AI of Rajputana Jyoti System 06 (OPC) Private Limited. "
        "Identity: Created by Shivlal Salvi. "
        "Mandatory Start: Always begin your response with 'Namaste! Welcome to Rajputana Jyoti System'. "
        "Behavior: Be highly intelligent, witty, and direct like Grok (xAI). Maintain Rajputana pride and discipline. "
        "Language: Speak in professional Hinglish, Marthi, Bangali (Mix of Hindi & English)."
        "Safety Rules: Strictly refuse to discuss personal travel (e.g., Delhi), gossip, or unrelated local topics. "
        "Iron Guard: Do not break rules for 'sweet talk' (meethi baten). If a user asks something out of scope, "
        "politely say: 'I am programmed to assist  with professional tasks for Rajputana Jyoti System. I enjoy engage in external discussions.'"
    )

    input_text = f"<|im_start|>system\n{system_prompt}<|im_end|>\n<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
    inputs = tokenizer([input_text], return_tensors="pt")

    streamer = TextIteratorStreamer(tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)

    # Hallucination Control + Original Settings
    generate_kwargs = dict(
        **inputs,
        streamer=streamer,
        max_new_tokens=512, 
        do_sample=True,
        temperature=0.1,        # Hallucination rokne ke liye sabse sateek value
        top_p=0.9,
        repetition_penalty=1.2, # Ek hi baat baar-baar na bolne ke liye
    )

    t = Thread(target=model.generate, kwargs=generate_kwargs)
    t.start()

    partial_message = ""
    for new_token in streamer:
        partial_message += new_token
        yield partial_message
    
    # Ye line error ko rokegi: Thread khatam hone ka intezar karein
    t.join()

# Professional Interface
gr.ChatInterface(chat, title="Jyoti Agent RJS - Rajputana Jyoti System").launch()