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

# Model ID
model_id = "mistralai/Mistral-7B-Instruct-v0.1"

# Load tokenizer & model
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype="auto",
    device_map=None  # ensures CPU runs safely
)

# Create pipeline
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

def chat(prompt):
    outputs = pipe(
        prompt,
        max_new_tokens=200,
        do_sample=True,
        temperature=0.7,
        top_p=0.9
    )
    return outputs[0]["generated_text"]

# Gradio interface
iface = gr.Interface(fn=chat, inputs="text", outputs="text", title="Super AI Brain")
iface.launch()