|
from transformers import AutoModelForCausalLM, AutoTokenizer |
|
import transformers |
|
import torch |
|
import gradio as gr |
|
|
|
|
|
torch.set_default_device("cuda") |
|
|
|
device = 'cuda' if torch.cuda.is_available() else 'cpu' |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
model = AutoModelForCausalLM.from_pretrained("deepapaikar/katzbot-phi2", |
|
torch_dtype=torch.float16, |
|
device_map="auto", |
|
trust_remote_code=True) |
|
tokenizer = AutoTokenizer.from_pretrained("deepapaikar/katzbot-phi2", trust_remote_code=True) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def predict_answer(question, token=25): |
|
messages = [{"role": "user", "content": f"{question}"}] |
|
|
|
|
|
prompt_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) |
|
|
|
|
|
inputs = tokenizer(prompt_text, return_tensors="pt", padding=True, truncation=True) |
|
|
|
|
|
inputs = {k: v.to(device) for k, v in inputs.items()} |
|
|
|
|
|
model.eval() |
|
model.to(device) |
|
|
|
|
|
output_sequences = model.generate( |
|
input_ids=inputs['input_ids'], |
|
attention_mask=inputs['attention_mask'], |
|
max_length=token + inputs['input_ids'].shape[-1], |
|
do_sample=True, |
|
temperature=0.7, |
|
top_k=50, |
|
top_p=0.95 |
|
) |
|
|
|
|
|
output_text = tokenizer.decode(output_sequences[0], skip_special_tokens=True) |
|
|
|
return output_text |
|
|
|
|
|
|
|
def gradio_predict(question, token): |
|
answer = predict_answer(question, token) |
|
return answer |
|
|
|
|
|
iface = gr.Interface( |
|
fn=gradio_predict, |
|
inputs=[gr.Textbox(label="Question", placeholder="e.g. Where is Yeshiva University located?", scale=4), |
|
gr.Slider(2, 100, value=25, label="Token Count", info="Choose between 2 and 100")], |
|
outputs=gr.TextArea(label="Answer"), |
|
title="KatzBot", |
|
description="Phi2-trial1", |
|
) |
|
|
|
|
|
iface.queue().launch(debug=True) |
|
|