Upload 3 files

app.py

@@ -0,0 +1,16 @@
+import gradio as gr
+from llama_generate import run
+
+
+def greet(query):
+    results = run(query)
+    return results
+
+
+sample_list = [
+    "Who is Lihu Chen?.",
+    "Who is Gaël Varoquaux?"
+]
+
+iface = gr.Interface(fn=greet, inputs="text", outputs="text", examples=sample_list, cache_examples=False)
+iface.launch()

llama_generate.py

@@ -0,0 +1,162 @@
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+from nltk.tokenize import sent_tokenize
+
+torch.device('cuda' if torch.cuda.is_available() else 'cpu') # the device to load the model onto
+model_name_or_path = "TheBloke/Llama-2-7b-Chat-GPTQ"
+
+model = AutoModelForCausalLM.from_pretrained(model_name_or_path,
+                                             #torch_dtype=torch.float16,
+                                             device_map="auto",
+                                             trust_remote_code=False,
+                                             revision="main")
+
+
+
+tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=True)
+tokenizer.pad_token = tokenizer.unk_token
+# tokenizer.add_special_tokens({'pad_token': '[PAD]'})
+# model.resize_token_embeddings(len(tokenizer))
+
+def clean(result):
+    special_token = ['<s>', '</s>', '<unk>']
+
+    result = result.split("[/INST]")[-1].strip()
+
+    # context = "[INST] {a} [/INST]".format(a=content)
+    #result = result.replace(context, '')
+    for token in special_token:
+        result = result.replace(token, '').strip()
+    return result.strip()
+
+def single_generate(query):
+
+    messages = [
+        {"role": "user", "content": query},
+    ]
+
+    encodeds = tokenizer.apply_chat_template(messages, return_tensors="pt")
+    #print(encodeds)
+
+    model_inputs = encodeds.to(device)
+    model.to(device)
+
+    generated_ids = model.generate(model_inputs, max_new_tokens=100,  do_sample=True, temperature=1.0)
+    decoded = tokenizer.batch_decode(generated_ids)
+    results = list()
+    for index, result in enumerate(decoded):
+        #print(result)
+        result = clean(result)
+        #print('query = ', query, ' result = ', result)
+        results.append(result)
+    return results
+
+
+def prepare_input(contents):
+    temp = list()
+    for content in contents:
+        messages = [
+            {"role": "user", "content": content}
+        ]
+        #print('messages = ', messages)
+        encodeds = tokenizer.apply_chat_template(messages, return_tensors="pt", max_length=100, padding=True)
+        # print(encodeds.size())
+        # print(encodeds)
+        temp.append(encodeds[0])
+
+    batch_encoded = torch.stack(temp)
+    return batch_encoded
+
+def batch_generate(queries):
+    model_inputs = prepare_input(queries).to(device)
+    model.to(device)
+
+    generated_ids = model.generate(model_inputs, max_new_tokens=100, do_sample=True, temperature=1.0)
+    decoded = tokenizer.batch_decode(generated_ids)
+    results = list()
+    for index, result in enumerate(decoded):
+        query = queries[index]
+        result = clean(result)
+        #print('query = ', query, ' result = ', result)
+        results.append(result)
+    return results
+
+
+def get_yes_or_no(result):
+    if 'yes' in str.lower(result)[:5]:return 'Yes'
+    if 'no' in str.lower(result)[:5]:return 'No'
+    return 'N/A'
+
+
+def check_score(context, sentences):
+    score_mapping = {'Yes':1.0, 'No':0.0}
+    template = """
+        Context: {a}
+        Sentence: {b}
+        Is the sentence supported by the context above?
+        Answer Yes or No (Don't give explanations):
+    """
+    scores, results = list(), list()
+    for sentence in sentences:
+        content = template.format(a=context.strip().replace('/n', ''), b=sentence.strip().replace('/n', ''))
+        result = single_generate(content)[0]
+        #result = clean(result, context)
+        #print('results', results)
+        results.append(result)
+
+    results = [get_yes_or_no(r) for r in results]
+    scores = [score_mapping.get(result, 0.5) for result in results]
+
+    # for sent, score in zip(sentences, scores):
+    #     print(sent.strip(), score)
+        #result_string += sent + ' ({a})'.format(a=score)
+
+    return scores
+
+
+def sample_answer(query, num):
+    answers = list()
+    for _ in range(num):
+        answer = single_generate(query)
+        answers.append(answer[0])
+    return answers
+
+
+def run(query, sample_size=5):
+    sampled = sample_answer(query, sample_size+1)
+    answer = sampled[0]
+    proofs = sampled[1:]
+    sentences = sent_tokenize(answer)
+
+    all_scores = list()
+    for proof in proofs:
+        scores = check_score(proof, sentences)
+        all_scores.append(scores)
+
+    final_content = ''
+    avg_confidence = list()
+    for index, scores in enumerate(zip(*all_scores)):
+        sentence_confidence = sum(scores) / len(scores)
+        avg_confidence.append(sentence_confidence)
+        final_content += sentences[index].strip() + ' ({a}) '.format(a=sentence_confidence)
+    avg_confidence = sum(avg_confidence) / len(avg_confidence)
+    final_content += '\nThe confidence score of this answer is {a}'.format(a=avg_confidence)
+    return final_content
+
+if __name__ == '__main__':
+    # result = sample_answer(query="Who is Lihu Chen?", num=5)
+    # print(result)
+    #batch_generate(["Who is Lihu Chen?", "Who is Lihu Chen?"])
+
+    # context = """
+    # Lihu Chen is an American writer and artist who works in comics. They received their degree in psychology from California State University, Fullerton and have worked on titles such as "The Gathering Storm" and "Heartthrob".
+    # """
+    # sentences = sent_tokenize("""
+    # Lihu Chen is an American writer and artist who works in comics. They received their degree in psychology from California State University, Fullerton and have worked on titles such as "The Gathering Storm" and "Heartthrob".
+    # """)
+    # result = check_score(context, sentences)
+    # print(result)
+    # result = """
+
+    answer = run(query='WHo is Lihu Chen?', sample_size=10)
+    print(answer)

requirements.txt

@@ -0,0 +1,3 @@
+nltk
+torch==2.1.0
+transformers==4.35.0.dev0