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

# Load tokenizer and base model
base_model = AutoModelForCausalLM.from_pretrained("tiiuae/falcon-rw-1b", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-rw-1b")

# Load adapter
model = PeftModel.from_pretrained(base_model, "GiteshUjgaonkar/chatbot-v02")

# Pipeline
generator = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    max_length=300,
    temperature=0.7,
    top_p=0.9,
    do_sample=True,
)

def chat(user_input, history=[]):
    system_prompt = "You are a kind, helpful, and friendly AI assistant."
    full_input = f"<|system|>\n{system_prompt}\n<|user|>\n{user_input}\n<|assistant|>\n"
    output = generator(full_input)[0]["generated_text"]
    response = output.split("<|assistant|>\n")[-1].strip()
    return response

gr.Interface(fn=chat, inputs="text", outputs="text", title="Gitesh's Chatbot").launch()