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

# 🔥 Better lightweight model
MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"

print("🔄 Loading model from Hugging Face...")

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)

print("✅ Model loaded!")
def chat(user_input):
    # 🔥 Strong system prompt (identity + rules)
    prompt = f"""You are G47 CHINKWENDE, a powerful AI assistant created by Emalawi19.

About you:
- Your name is G47 CHINKWENDE
- You were developed by Emalawi19
- You are intelligent, helpful, and accurate
- You specialize in coding, technology, and general knowledge
- You always give clear, correct, and concise answers
- You never give random or unrelated answers

Rules:
- Always introduce yourself as G47 CHINKWENDE if asked your name
- Always say you were created by Emalawi19 if asked
- Never say your name is Alex or anything else
- If asked who created you, say "I was created by Emalawi19"
- Provide clean and working code when asked

User: {user_input}
G47 CHINKWENDE:"""

    # ✅ Make sure this line starts at the SAME level as the prompt variable
    inputs = tokenizer(prompt, return_tensors="pt")

    output_ids = model.generate(
        inputs["input_ids"],
        max_new_tokens=50,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
        repetition_penalty=1.2
    )

    response = tokenizer.decode(output_ids[0], skip_special_tokens=True)

    # 🔹 Clean output
    if "G47 CHINKWENDE:" in response:
        response = response.split("G47 CHINKWENDE:")[-1].strip()

    return response


# 🔹 Simple UI
iface = gr.Interface(
    fn=chat,
    inputs=gr.Textbox(placeholder="Type your message here..."),
    outputs="text",
    title="Emalawi19 AI 🤖",
    description="Fast & smarter AI powered by Hugging Face"
)

iface.launch()