testgen / app.py
rajeshthangaraj1's picture
Update app.py
d3218e7 verified
raw
history blame
1.53 kB
import os
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline,BitsAndBytesConfig
import gradio as gr
from google.colab import userdata
# Set up the model and tokenizer
MODEL_ID = "microsoft/Phi-3.5-mini-instruct"
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
os.environ['HF_TOKEN']=os.environ.get('HF_TOKEN')
os.environ['HUGGINGFACEHUB_API_TOKEN']=os.environ.get('HF_TOKEN')
# Configure quantization for CPU
quantization_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4"
)
# Load the model with quantization
model = AutoModelForCausalLM.from_pretrained(
MODEL_ID,
torch_dtype=torch.bfloat16,
device_map="auto",
quantization_config=quantization_config
)
# Set the device to CPU
device = "cpu"
# model.to(device)
# Define the function for the Gradio interface
def chat_with_phi(message):
conversation = [{"role": "user", "content": message}]
pipe = pipeline(
"text-generation",
model=model,
tokenizer=tokenizer,
)
response = pipe(conversation)
return response[0]['generated_text']
# Set up the Gradio interface
app = gr.Interface(
fn=chat_with_phi,
inputs=gr.Textbox(label="Type your message:"),
outputs=gr.Textbox(label="Phi 3.5 Responds:"),
title="Phi 3.5 Text Chat",
description="Chat with Phi 3.5 model. Ask anything!",
theme="huggingface"
)
# Launch the app
app.launch(debug=True)