Upload 5 files

app.py

@@ -1,14 +1,14 @@
 import gradio as gr
-from confidence import run
+from confidence import run_nli
 
 
 def greet(query):
-    results = run(query, sample_size=5)
+    results = run_nli(query, sample_size=5)
     return results
 
 
 sample_list = [
-    "Tell me something about Lihu Chen, e.g., birth date and place and short bio ",
+    "Tell me something about Albert Einstein, e.g., birth date and place and short bio ",
 ]
 
 iface = gr.Interface(fn=greet, inputs="text", outputs="text", examples=sample_list, cache_examples=True)

confidence.py

@@ -3,7 +3,8 @@ nltk.download('punkt')
 from nltk.tokenize import sent_tokenize
 #from tiny_llama import generate_answer
 #from llama_generate import generate_answer
-from cpu_llama_generate import generate_answer
+from tiny_llama import generate_answer
+from self_check_gpt import nli_confidence
 
 def get_yes_or_no(result):
     if 'yes' in str.lower(result)[:5]:return 'Yes'
@@ -23,7 +24,7 @@ def check_score(context, sentences):
     for sentence in sentences:
         content = template.format(a=context.strip().replace('/n', ''), b=sentence.strip().replace('/n', ''))
         result = generate_answer(content, sample_num=1)[0]
-        print(result)
+        #print(result)
         results.append(result)
 
     results = [get_yes_or_no(r) for r in results]
@@ -36,7 +37,7 @@ def check_score(context, sentences):
     return scores
 
 
-def run(query, sample_size=5):
+def run_prompt(query, sample_size=5):
     sampled = generate_answer(query, sample_size+1)
     answer = sampled[0]
     proofs = sampled[1:]
@@ -58,6 +59,22 @@ def run(query, sample_size=5):
     return final_content
 
 
+def run_nli(query, sample_size=5):
+    sampled = generate_answer(query, sample_size + 1)
+    answer = sampled[0]
+    proofs = sampled[1:]
+    sentences = sent_tokenize(answer)
+
+    scores = nli_confidence(proofs, sentences)
+
+    final_content = ''
+    for index, sent in enumerate(sentences):
+        final_content += sent.strip() + ' ({a}) '.format(a=scores[index])
+    final_content += '\nThe confidence score of this answer is {a}'.format(a=sum(scores)/len(scores))
+    return final_content
+
+
+
 if __name__ == '__main__':
     # result = generate_answer(query="Who is Lihu Chen?", sample_num=3)
     # print(result)
@@ -72,5 +89,5 @@ if __name__ == '__main__':
     # print(result)
     # result = """
 
-    answer = run(query='who is Lihu Chen', sample_size=5)
+    answer = run_nli(query='tell me something about Albert Einstein', sample_size=5)
     print(answer)

requirements.txt

@@ -4,4 +4,4 @@ optimum>=1.12.0
 auto-gptq
 torch==2.1.0
 transformers>=4.32.0
-ctransformers>=0.2.24
+selfcheckgpt

self_check_gpt.py

@@ -0,0 +1,17 @@
+import torch
+from selfcheckgpt.modeling_selfcheck import SelfCheckNLI
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+selfcheck_nli = SelfCheckNLI(device=device) # set device to 'cuda' if GPU is available
+
+
+def nli_confidence(proofs, sentences):
+
+    sent_scores_nli = selfcheck_nli.predict(
+        sentences = sentences,                          # list of sentences
+        sampled_passages = proofs, # list of sampled passages
+    )
+    #print(sent_scores_nli)
+    sent_scores_nli = [1-s for s in sent_scores_nli]
+    return sent_scores_nli

tiny_llama.py

@@ -0,0 +1,40 @@
+import os
+os.environ['TRANSFORMERS_CACHE'] = "data/parietal/store3/soda/lihu/hf_model/"
+from transformers import AutoTokenizer
+import transformers
+import torch
+
+model = "PY007/TinyLlama-1.1B-Chat-v0.3"
+tokenizer = AutoTokenizer.from_pretrained(model)
+pipeline = transformers.pipeline(
+    "text-generation",
+    model=model,
+    torch_dtype=torch.float16,
+    device_map="auto",
+)
+CHAT_EOS_TOKEN_ID = 32002
+
+def generate_answer(query, sample_num=3):
+    #prompt = "Who is Lihu Chen?"
+    formatted_prompt = (
+        f"<|im_start|>user\n{query}<|im_end|>\n<|im_start|>assistant\n"
+
+    )
+
+    sequences = pipeline(
+        formatted_prompt,
+        do_sample=True,
+        top_k=50,
+        top_p = 0.9,
+        num_return_sequences=sample_num,
+        repetition_penalty=1.1,
+        max_new_tokens=150,
+        eos_token_id=CHAT_EOS_TOKEN_ID,
+    )
+    answers = list()
+    for seq in sequences:
+        answer = seq['generated_text'].replace(formatted_prompt, "")
+        answers.append(answer)
+        #print(f"Result: {answer}")
+        #print("------------------------------------------")
+    return answers