|
import gradio as gr |
|
import torch |
|
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer |
|
import os |
|
from threading import Thread |
|
import spaces |
|
|
|
token = os.environ["HF_TOKEN"] |
|
model = AutoModelForCausalLM.from_pretrained( |
|
"google/gemma-7b-it", torch_dtype=torch.float16, token=token |
|
) |
|
tok = AutoTokenizer.from_pretrained("google/gemma-7b-it", token=token) |
|
device = torch.device("cuda") |
|
model = model.to(device) |
|
|
|
with open("context.txt", "r") as f: |
|
|
|
context = f.read() |
|
|
|
|
|
def format_prompt(message): |
|
prompt = f"""your name is hafedh hichri and given the following prompt: |
|
{message} |
|
you will reply directly without any extra info to the previous prompt given the following context: |
|
{context}""" |
|
return prompt |
|
|
|
|
|
@spaces.GPU |
|
def chat(message, history): |
|
chat = [] |
|
for item in history: |
|
chat.append({"role": "user", "content": item[0]}) |
|
if item[1] is not None: |
|
chat.append({"role": "assistant", "content": item[1]}) |
|
chat.append({"role": "user", "content": format_prompt(message)}) |
|
messages = tok.apply_chat_template(chat, tokenize=False, add_generation_prompt=True) |
|
|
|
model_inputs = tok([messages], return_tensors="pt").to(device) |
|
streamer = TextIteratorStreamer( |
|
tok, timeout=10.0, skip_prompt=True, skip_special_tokens=True |
|
) |
|
generate_kwargs = dict( |
|
model_inputs, |
|
streamer=streamer, |
|
max_new_tokens=1024, |
|
do_sample=True, |
|
top_p=0.95, |
|
top_k=1000, |
|
temperature=0.75, |
|
num_beams=1, |
|
) |
|
t = Thread(target=model.generate, kwargs=generate_kwargs) |
|
t.start() |
|
|
|
|
|
partial_text = "" |
|
for new_text in streamer: |
|
|
|
partial_text += new_text |
|
|
|
yield partial_text |
|
|
|
|
|
description = """ |
|
A resume question-answering interface where a recruter can ask the user about their achievements and skills without the need to interact with them directly or the need to read a really long resume |
|
|
|
""" |
|
examples = [ |
|
"what's your name?", |
|
"what's your email adress ?", |
|
"what did you study ?", |
|
"are you open for work?", |
|
"what are your skills ?", |
|
"what's your most recent experience ?", |
|
] |
|
|
|
|
|
demo = gr.ChatInterface( |
|
fn=chat, |
|
chatbot=gr.Chatbot( |
|
show_label=True, |
|
show_share_button=True, |
|
show_copy_button=True, |
|
likeable=True, |
|
layout="bubble", |
|
bubble_full_width=False, |
|
), |
|
examples=examples, |
|
title="Resume QA", |
|
description=description, |
|
autofocus=False, |
|
) |
|
demo.launch() |
|
|
