|
import gradio as gr |
|
import torch |
|
from transformers import AutoTokenizer |
|
|
|
class Pipline: |
|
def __init__(self, model, tokenizer, device='cpu'): |
|
self.device = device |
|
self.model = model.to(self.device) |
|
self.tokenizer = tokenizer |
|
self.pre_prompt = "\n\nYou are a AI assistant who helps the user to solve their issue\n\n" |
|
|
|
@torch.no_grad() |
|
def respond(self, Instruction=None, input=None, temperature=0.8, max_length=200, do_sample=True, top_k=0, top_p=0.9, repetition_penalty=1.0, num_return_sequences=1, num_beams=1, early_stopping=False, use_cache=True, **generate_kwargs): |
|
if not Instruction and not input: |
|
raise ValueError("Either Instruction or input must be passed.") |
|
query = f"""{self.pre_prompt} |
|
Instruction: {Instruction if Instruction else ""} |
|
Input: {input if input else ""} |
|
Output:""" |
|
inp_tokens_l = self.tokenizer(query, return_tensors='pt').input_ids |
|
inp_tokens = inp_tokens_l.to(self.device) |
|
out_tokens = self.model.generate(inp_tokens, max_length=max_length, temperature=temperature, do_sample=do_sample, top_k=top_k, top_p=top_p, repetition_penalty=repetition_penalty, num_return_sequences=num_return_sequences, num_beams=num_beams, early_stopping=early_stopping, use_cache=use_cache, **generate_kwargs) |
|
out_text = self.tokenizer.batch_decode(out_tokens, skip_special_tokens=True) |
|
|
|
return out_text |
|
|
|
tokenizer = AutoTokenizer.from_pretrained("facebook/opt-125M") |
|
model = torch.load('./model-cpu.pkl') |
|
|
|
pipe = Pipline(model=model, tokenizer=tokenizer, device='cpu') |
|
|
|
input_components = [ |
|
gr.inputs.Textbox(label='Instruction', placeholder='Enter instruction...'), |
|
gr.inputs.Textbox(label='Input', placeholder='Enter input...'), |
|
] |
|
|
|
output_components = [ |
|
gr.outputs.Textbox(label='Output'), |
|
] |
|
|
|
def chatbot_response(Instruction, input, max_length, temperature): |
|
output = pipe.respond( |
|
Instruction=Instruction, |
|
input=input, |
|
max_length=int(max_length), |
|
temperature=float(temperature), |
|
) |
|
return output[0] |
|
|
|
interface = gr.Interface( |
|
fn=chatbot_response, |
|
inputs=input_components + [ |
|
gr.inputs.Slider( |
|
label='Max Length', |
|
minimum=10, |
|
maximum=500, |
|
step=10, |
|
default=200, |
|
), |
|
gr.inputs.Slider( |
|
label='Temperature', |
|
minimum=0.1, |
|
maximum=1.0, |
|
step=0.1, |
|
default=0.8, |
|
), |
|
], |
|
outputs=output_components, |
|
title='ChatOPT', |
|
description='Type in an instruction and input, and get a response from the model', |
|
) |
|
|
|
interface.launch() |
|
|
