API fix

app.py

@@ -6,11 +6,12 @@ from datetime import datetime
 import pandas as pd
 
 from utils import generate, send_to_s3
-from models import get_tinyllama, response_tinyllama
+from models import get_tinyllama, response_tinyllama, response_qwen2ins1b
 from constants import css, js_code, js_light
 
 MERA_table = None
-TINY_LLAMA = None
+TINYLLAMA = None
+QWEN2INS1B = None
 
 S3_SESSION = None
 
@@ -28,6 +29,14 @@ def tiny_gen(content, chat_history):
     send_to_s3(res, f'protobench/tiny_{str(datetime.now()).replace(" ", "_")}.json', S3_SESSION)
     return '', chat_history
 
+def qwen_gen(content, chat_history):
+    chat_history.append([content])
+    res = response_qwen2ins1b(QWEN2INS1B, chat_history)
+    chat_history[-1].append(res)
+    send_to_s3(res, f'protobench/tiny_{str(datetime.now()).replace(" ", "_")}.json', S3_SESSION)
+    return '', chat_history
+
+
 def tab_arena():
     with gradio.Row():
         with gradio.Column():
@@ -51,7 +60,7 @@ def tab_arena():
     #     return "", chat_history
 
     msg.submit(giga_gen, [msg, chatbot_left], [msg, chatbot_left])
-    msg.submit(tiny_gen, [msg, chatbot_right], [msg, chatbot_right])
+    msg.submit(qwen_gen, [msg, chatbot_right], [msg, chatbot_right])
 
         # with gradio.Column():
         #     gradio.ChatInterface(

models.py

@@ -5,6 +5,10 @@ def get_tinyllama():
     tinyllama = pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", torch_dtype=torch.float16, device_map="auto")
     return tinyllama
 
+def get_qwen2ins1b():
+    tinyllama = pipeline("text-generation", model="Qwen/Qwen2-1.5B-Instruct", torch_dtype=torch.float16, device_map="auto")
+    return tinyllama
+
 def response_tinyllama(
         model=None,
         messages=None
@@ -24,4 +28,25 @@ def response_tinyllama(
     prompt = model.tokenizer.apply_chat_template(messages_dict, tokenize=False, add_generation_prompt=True)
     outputs = model(prompt, max_new_tokens=64, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
 
-    return outputs[0]['generated_text'].split('<|assistant|>')[1].strip()
+    return outputs[0]['generated_text'].split('<|assistant|>')[1].strip()
+
+def response_qwen2ins1b(
+        model=None,
+        messages=None
+        ):
+    
+    messages_dict = [
+        {
+            "role": "system",
+            "content": "You are a friendly and helpful chatbot",
+        }
+    ]
+    for step in messages:
+        messages_dict.append({'role': 'user', 'content': step[0]})
+        if len(step) >= 2:
+            messages_dict.append({'role': 'assistant', 'content': step[1]})
+
+    prompt = model.tokenizer.apply_chat_template(messages_dict, tokenize=False, add_generation_prompt=True)
+    outputs = model(prompt, max_new_tokens=64, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
+
+    return outputs[0]['generated_text'] #.split('<|assistant|>')[1].strip()