app.py

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

# Load your custom model
model_name = "ShabanEjupi/Chatbot-i"

try:
    # Option 1: Load directly using pipeline (simpler)
    chatbot = pipeline(
        "text2text-generation",
        model=model_name,
        tokenizer=model_name,
        device=-1  # Use CPU
    )
    
    # Test the pipeline
    test_output = chatbot("Hello!", max_length=50)
    print("✓ Model loaded successfully")
    print("Test output:", test_output)
    
except Exception as e:
    print(f"Error loading with pipeline: {e}")
    print("Trying manual loading...")
    
    try:
        # Option 2: Manual loading if pipeline fails
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
        print("✓ Model loaded manually")
    except Exception as e:
        raise gr.Error(f"Failed to load model: {str(e)}")

def respond(message):
    try:
        # Format prompt for FLAN-T5
        prompt = f"User: {message}\nAssistant:"
        
        # Use whichever loading method worked
        if 'chatbot' in globals():
            response = chatbot(prompt, max_length=100)
            return response[0]['generated_text']
        else:
            inputs = tokenizer(prompt, return_tensors="pt")
            outputs = model.generate(**inputs, max_length=100)
            return tokenizer.decode(outputs[0], skip_special_tokens=True)
            
    except Exception as e:
        return f"Error generating response: {str(e)}"

# Create Gradio interface
demo = gr.Interface(
    fn=respond,
    inputs=gr.Textbox(lines=2, placeholder="Type your message here..."),
    outputs="text",
    title="FLAN-T5 Chatbot",
    examples=["Hello!", "What's AI?", "Explain quantum computing"],
    description="Chatbot using your custom FLAN-T5 model"
)

demo.launch()