app.py

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

# Model configuration
MODEL_NAME = "kuyesu22/Llama-3.2-3B-Instruct-Sunbird-Dialogue.v1"


# Quantization configuration for efficient model loading
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=True)

# Load model with quantization and device mapping
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    quantization_config=quantization_config,
    device_map="auto"
)

# Initialize pipeline
pipe = pipeline(
    task="text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=128,
    return_full_text=False,
    # device=0  # Ensure it runs on the GPU if available
)

# Function to generate predictions based on the question
def generate_response(question):
    # Construct a prompt based only on the user's question
    prompt = f"""You are an AI assistant specialized in answering questions related to prenatal, can you respond appropriately to the user only in Runyakole a language spoken by people in western Uganda. 
    Please provide a detailed response in Runyankore for the following question:
    Question: {question}
    Answer:"""
    
    # Generate answer using the pipeline
    outputs = pipe(prompt)
    return outputs[0]["generated_text"]

# Gradio Interface
def gradio_interface(question):
    if not question.strip():
        return "Please enter a valid question."
    response = generate_response(question)
    return response

# Define Gradio inputs and outputs
gr_interface = gr.Interface(
    fn=gradio_interface,
    inputs=gr.components.Textbox(label="Enter your question", placeholder="e.g., Amakye maama genda kusula mu biseera bitya mu kucuma?"),
    outputs=gr.components.Textbox(label="Generated Answer"),
    title="Dialogue of Delivery  - Runyankore Q&A",
    description="Ask any question related to prenatal care in Runyankore, and get an AI-powered answer.",
    theme="default",
    allow_flagging="never",
)

# Launch Gradio app
if __name__ == "__main__":
    gr_interface.launch()