Arena test

app.py

@@ -3,21 +3,33 @@ import argparse
 import os
 
 from utils import generate
+from models import get_tiny_llama, response_tiny_llama
 from constants import css, js_code, js_light
 
 MERA_table = None
+TINY_LLAMA = get_tiny_llama()
 
-def gen(content):
+def giga_gen(content):
     res = generate(content,'auth_token.json')
     return res
 
+def tiny_gen(content):
+    res = response_tiny_llama(TINY_LLAMA, content)
+    return res
+
 def tab_arena():
-    gradio.Interface(fn=gen, inputs="text", outputs="text") # arena = 
+    with gradio.Row():
+        with gradio.Column():
+            gradio.Interface(fn=giga_gen, inputs="text", outputs="text", allow_flagging=False, title='Giga') # arena = 
+        with gradio.Column():
+            gradio.Interface(fn=tiny_gen, inputs="text", outputs="text", allow_flagging=False, title='TinyLlama') # arena = 
+        
     # arena.launch()
 
 with open("test.md", "r") as f:
     TEST_MD = f.read()
 
+available_models = ["GigaChat", ""] # list(model_info.keys()) 
 
 def build_demo():
     # global original_dfs, available_models, gpt4t_dfs, haiku_dfs, llama_dfs

models.py

@@ -0,0 +1,23 @@
+import torch
+from transformers import pipeline
+
+def get_tiny_llama():
+    pipe = pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", torch_dtype=torch.float16, device_map="auto")
+    return pipe
+
+def response_tiny_llama(
+        pipe=None,
+        content="How many helicopters can a human eat in one sitting?"
+        ):
+    # We use the tokenizer's chat template to format each message - see https://huggingface.co/docs/transformers/main/en/chat_templating
+    messages = [
+        {
+            "role": "system",
+            "content": "You are a friendly chatbot who always responds in the style of a pirate",
+        },
+        {"role": "user", "content": content},
+    ]
+    prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    outputs = pipe(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
+
+    return outputs[0]['generated_text']