import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load Dolly 2.0 model and tokenizer
model_name = "databricks/dolly-v2-3b"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Define a function to generate responses
def dolly_response(prompt):
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(
        inputs["input_ids"],
        max_length=50,
        num_beams=3,
        no_repeat_ngram_size=2,
        early_stopping=True
    )
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Create Gradio interface
gr.Interface(
    fn=dolly_response,
    inputs="text",
    outputs="text",
    title="Dolly 2.0 Chatbot"
).launch()