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

# Model ID
model_id = "large-traversaal/Alif-1.0-8B-Instruct"

# Load tokenizer and model (CPU-friendly)
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cpu")  # Changed to CPU

# Create text generation pipeline
chatbot = pipeline("text-generation", model=model, tokenizer=tokenizer, device="cpu")  # Ensuring CPU use

# Function to generate responses
def chat(message):
    response = chatbot(message, max_new_tokens=100, do_sample=True, temperature=0.3)
    return response[0]["generated_text"]

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# 🤖 Alif Chatbot - Urdu Language AI Model")
    user_input = gr.Textbox(label="User Input", placeholder="اپنا سوال یہاں لکھیں...")
    submit_btn = gr.Button("Send")
    bot_response = gr.Textbox(label="AI Response")

    submit_btn.click(fn=chat, inputs=user_input, outputs=bot_response)

# Launch the app
if __name__ == "__main__":
    demo.launch()