Replace front

app.py

@@ -28,25 +28,13 @@ async def handle(system_input: str, user_input: str):
 with gr.Blocks() as demo:
     system_input = gr.Textbox(label='System Input', value='You are AI assistant', lines=2)
     with gr.Row():
-        with gr.Column():
-            gpt = gr.Textbox(label='gpt-2', lines=4, interactive=False)
-            gr.Text("OpenAI\n14 February 2019")
-        with gr.Column():
-            t5 = gr.Textbox(label='t5', lines=4, interactive=False)
-            gr.Text("Google\n12 Dec 2019")
-        with gr.Column():
-            bloom = gr.Textbox(label='bloom [GPU]', lines=4, interactive=False)
-            gr.Text('Big Science\n11 Jul 2022')
+        gpt = gr.Textbox(label='gpt-2', lines=4, interactive=False, info='OpenAI\n14 February 2019')
+        t5 = gr.Textbox(label='t5', lines=4, interactive=False, info='Google\n12 Dec 2019')
+        bloom = gr.Textbox(label='bloom [GPU]', lines=4, interactive=False, info='Big Science\n11 Jul 2022')
     with gr.Row():
-        with gr.Column():
-            llama2 = gr.Textbox(label='llama-2', lines=4, interactive=False)
-            gr.Text("MetaAI\n18 Jul 2023")
-        with gr.Column():
-            mistral = gr.Textbox(label='mistral-v01', lines=4, interactive=False)
-            gr.Text("MistralAI\n20 Sep 2023")
-        with gr.Column():
-            llama3 = gr.Textbox(label='llama-3.1', lines=4, interactive=False)
-            gr.Text('MetaAI\n18 Jul 2024')
+        llama2 = gr.Textbox(label='llama-2', lines=4, interactive=False, info='MetaAI\n18 Jul 2023')
+        mistral = gr.Textbox(label='mistral-v01', lines=4, interactive=False, info='MistralAI\n20 Sep 2023')
+        llama3 = gr.Textbox(label='llama-3.1', lines=4, interactive=False, info='MetaAI\n18 Jul 2024')
 
     user_input = gr.Textbox(label='User Input', lines=2, value='Calculate expression: 7-3=')
     gen_button = gr.Button('Generate')

generators.py

@@ -11,7 +11,7 @@ from huggingface_hub import InferenceClient
 import random
 import torch
 from huggingface_hub import AsyncInferenceClient
-from transformers import LlamaTokenizer, LlamaForCausalLM, AutoTokenizer
+from transformers import LlamaTokenizer, LlamaForCausalLM, AutoTokenizer, AutoModelForCausalLM
 
 
 async def query_llm(payload, model_name):
@@ -85,7 +85,7 @@ def generate_openllama(system_input, user_input):
 def generate_bloom(system_input, user_input):
     model_path = 'bigscience/bloom-7b1'
     tokenizer = AutoTokenizer.from_pretrained(model_path)
-    model = LlamaForCausalLM.from_pretrained(
+    model = AutoModelForCausalLM.from_pretrained(
         model_path, torch_dtype=torch.float16, device_map='cuda',
     )
     input_text = f"{system_input}\n{user_input}"