import os
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    pipeline,
)
from transformers import BitsAndBytesConfig

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype="float16",
    bnb_4bit_use_double_quant=False,
)

model_name = "rukaiyah-indika-ai/rv-chatbot-2"

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config
)
    
import gradio as gr

def generate_response(prompt):
    inst = "You are a very helpful assistant providing solutions to road-related queries. Ensure you provide correct and relevant answers according to the IRC guidelines. If you don't know the answer to a question, please don't share false information."
    pipe = pipeline(task="text-generation", model=model, tokenizer=tokenizer, temperature=0.2, max_new_tokens=256)
    ranked_results = pipe(f"<s>[INST] {inst}{prompt} [/INST]")

    for result in ranked_results:
        response = result['generated_text']
        response = response.split("[/INST]", 1)[-1]
        response = response.replace("<s>", "")
        response = response.replace("</s>", "")
        return response

iface = gr.Interface(
    fn=generate_response,
    inputs="text",
    outputs="text",
    title="Road-GPT",
    description="Enter your query related to road management and get a response generated by our chatbot."
)

iface.launch(share=True)