switch to a text generation model

app.py

@@ -1,22 +1,94 @@
 import gradio as gr
-from transformers import pipeline
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+import os
+from threading import Thread
+import spaces
 
-qa_model = pipeline("question-answering", "timpal0l/mdeberta-v3-base-squad2")
+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)
 
-description="""
+with open("context.txt", "r") as f:
+    # read content from the resume
+    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)
+    # Tokenize the messages string
+    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()
+
+    # Initialize an empty string to store the generated text
+    partial_text = ""
+    for new_text in streamer:
+        # print(new_text)
+        partial_text += new_text
+        # Yield an empty string to cleanup the message textbox and the updated conversation history
+        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
 
-**LIMITATIONS:** the bot can only extract specific information and does not take into account multiple sentences at once.
 """
-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 ?"]
-
-with open("context.txt","r") as f:
-  # read content from the resume
-  context = f.read()
+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 ?",
+]
 
-def chat(question,history):
-  """chat with the QA pipeline"""
-  return qa_model(question = question, context = context)["answer"].strip()
 
-demo = gr.ChatInterface(fn=chat, examples=examples, title="Resume QA",description=description,autofocus=False)
+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()

requirements.txt

@@ -1,2 +1,3 @@
 transformers
-torch
+torch==2.2.0
+spaces