main app for topk responses API

app.py

@@ -0,0 +1,136 @@
+import os
+import urllib.request
+import gradio as gr
+from transformers import T5Tokenizer, T5ForConditionalGeneration
+import huggingface_hub
+import re
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+import time
+import transformers
+import requests
+import globals
+from utility import *
+
+"""set up"""
+huggingface_hub.login(token=globals.HF_TOKEN)
+gemma_tokenizer = AutoTokenizer.from_pretrained(globals.gemma_2b_URL)
+gemma_model = AutoModelForCausalLM.from_pretrained(globals.gemma_2b_URL)
+
+falcon_tokenizer = AutoTokenizer.from_pretrained(globals.falcon_7b_URL, trust_remote_code=True, device_map=globals.device_map, offload_folder="offload")
+falcon_model = AutoModelForCausalLM.from_pretrained(globals.falcon_7b_URL, trust_remote_code=True,
+                                                    torch_dtype=torch.bfloat16, device_map=globals.device_map, offload_folder="offload")
+
+def get_model(model_typ):
+  if model_typ not in ["gemma", "falcon", "falcon_api", "simplet5_base", "simplet5_large"]:
+    raise ValueError('Invalid model type. Choose "gemma", "falcon", "falcon_api","simplet5_base", "simplet5_large".')
+  if model_typ=="gemma":
+    tokenizer = gemma_tokenizer
+    model = gemma_model
+    prefix = globals.gemma_PREFIX
+  elif model_typ=="falcon_api":
+    prefix = globals.falcon_PREFIX
+    model=None
+    tokenizer = None
+  elif model_typ=="falcon":
+    tokenizer = falcon_tokenizer
+    model = falcon_model
+    prefix = globals.falcon_PREFIX
+  elif model_typ in ["simplet5_base","simplet5_large"]:
+    prefix = globals.simplet5_PREFIX
+    URL = globals.simplet5_base_URL if model_typ=="simplet5_base" else globals.simplet5_large_URL
+    T5_MODEL_PATH = f"https://huggingface.co/{URL}/resolve/main/{globals.T5_FILE_NAME}"
+    fetch_model(T5_MODEL_PATH, globals.T5_FILE_NAME)
+    tokenizer = T5Tokenizer.from_pretrained(URL)
+    model = T5ForConditionalGeneration.from_pretrained(URL)
+  return model, tokenizer, prefix
+
+def topk_query(model_typ="gemma",prompt="She has a heart of gold",temperature=0.7,max_length=256):
+  if model_typ not in ["gemma","simplet5_base", "simplet5_large"]:
+    raise ValueError('Invalid model type. Choose "gemma", "simplet5_base", "simplet5_large".')
+  model, tokenizer, prefix = get_model(model_typ)
+
+  start_time = time.time()
+  input = prefix.replace("{fig}", prompt)
+  print(f"Input to model: \n{input}")
+
+  if model_typ in ["simplet5_base", "simplet5_large"]:
+    inputs = tokenizer(input, return_tensors="pt")
+    outputs = model.generate(
+        inputs["input_ids"],
+        temperature=temperature,
+        max_length=max_length,
+        num_beams=5,
+        num_return_sequences=5,  # Generate 5 responses
+        early_stopping=True
+    )
+
+    response = [tokenizer.decode(output, skip_special_tokens=True) for output in outputs]
+    answer = [response.replace(input, "").strip() for response in response]
+
+  elif model_typ=="gemma":
+    inputs = tokenizer(input, return_tensors="pt")
+    generate_ids = gemma_model.generate(
+        inputs.input_ids,
+        max_length=max_length,
+        do_sample=True,
+        top_k=50,
+        temperature=temperature,
+        num_return_sequences=5,
+        eos_token_id=gemma_tokenizer.eos_token_id
+    )
+    outputs = gemma_tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+    print(f"Model original output:{outputs}\n")
+    answer = [post_process(output,input).replace("\n", "") for output in outputs]
+
+  # TODO: falcon's outputs dont have much differences, not used in topk response
+  # elif model_typ=="falcon_api":
+  #   API_URL = "https://api-inference.huggingface.co/models/tiiuae/falcon-7b-instruct"
+  #   headers = {"Authorization": f"Bearer {access_token}"}
+  #   response = api_query(API_URL=API_URL, headers=headers, payload={
+  #       "inputs": input,
+  #       "parameters": {
+  #           "temperature": temperature,
+  #           "top_k": 50,
+  #           "num_return_sequences": 5
+  #       }
+  #   })
+  #   print(response)
+  #   answer = [post_process(item["generated_text"], input) for item in response]
+  else:
+    raise ValueError('Invalid model type. Choose "gemma", "simplet5_base", "simplet5_large".')
+
+  print(f"Time taken: {time.time()-start_time:.2f} seconds")
+  print(f"processed model output: {answer}")
+
+  return answer
+
+topk_iface = gr.Interface(
+    fn=topk_query,
+    inputs=[
+        gr.Dropdown(
+            choices=["gemma", "simplet5_base", "simplet5_large"],
+            label="Model Type",
+            value="gemma"
+        ),
+        gr.Textbox(label="Prompt", placeholder="Enter your prompt here"),
+        gr.Slider(minimum=0.1, maximum=1.0, step=0.1, value=0.7, label="Temperature"),
+        gr.Slider(minimum=50, maximum=512, step=10, value=256, label="Max Length")
+    ],
+    outputs=[
+        gr.Textbox(label="Response 1"),
+        gr.Textbox(label="Response 2"),
+        gr.Textbox(label="Response 3"),
+        gr.Textbox(label="Response 4"),
+        gr.Textbox(label="Response 5")
+    ],theme=gr.themes.Soft(),
+    title=globals.TITLE,
+    description="Generate multiple responses (top 5) based on input sentence, prefix, and temperature. Literal meanings/explanations are provided based on the input figurative sentence.",
+    examples=[
+           ["gemma", "Time flies when you're having fun",0.7],
+           ["simplet5_large", "She has a heart of gold",0.5],
+        ["gemma", "The sky is the limit",0.6]
+    ]
+)
+if __name__ == '__main__':
+  topk_iface.launch()

requirements.txt

@@ -0,0 +1,7 @@
+huggingface_hub>=0.23.2
+llama-cpp-python==0.1.62
+transformers==4.42.4
+sentence_transformers
+sentencepiece
+accelerate==0.32.1
+bitsandbytes==0.43.3