# app.py

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

# Load the model and tokenizer from Hugging Face
model_path = "Athagi/Agillm-v2"  # Model loaded from Hugging Face Hub
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(model_path)

# Create a text generation pipeline
chatbot = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1)

# Function to generate response
def chat_with_model(user_input):
    response = chatbot(user_input, max_length=200, do_sample=True, temperature=0.7)
    return response[0]['generated_text']

# Gradio interface
interface = gr.Interface(
    fn=chat_with_model,  # Function to generate the response
    inputs="text",  # Input: Text box for user input
    outputs="text",  # Output: Generated response
    title="Chat with Agillm-v2",  # Title of the app
    description="Type a message and interact with the Agillm-v2 model.",
    theme="huggingface"  # Optional: Use Hugging Face theme
)

# Launch the app
interface.launch()