roadGPT / app.py
rukaiyaaaah's picture
Update app.py
20a7c22 verified
raw
history blame contribute delete
No virus
1.46 kB
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)