add ability to return logprobs for lmtp models

Author: fjfricke

src/lmql/models/lmtp/lmtp_dcmodel.py

@@ -224,7 +224,7 @@ def make_logits(self, payload):
     async def singleton_result(self, token, score):
         yield {"token": token, "logprob": score, "top_logprobs": {token: score}}
 
-    async def generate(self, s, temperature, top_logprobs = 1, chunk_size=None, **kwargs):
+    async def generate(self, s, temperature, top_logprobs = 5, chunk_size=None, **kwargs):
         kwargs = {**self.model_args, **kwargs}
 
         # get token masks from interpreter

src/lmql/runtime/dclib/dclib_seq.py

@@ -441,7 +441,7 @@ def make_successors(self, next_tokens, next_token_scores, logits, user_data=None
         tokens = [t for t, s in zip(next_tokens, next_token_scores) if s > get_truncation_threshold()]
         scores = [s for s in next_token_scores if s > get_truncation_threshold()]
 
-        distribution_logprobs = [{k: v for k, v in logits.probs.items() if type(k) == str}]
+        distribution_logprobs = [{get_tokenizer().decode([k]): v for k, v in logits.probs.items() if type(k) == int and v > -10}]
         if len(distribution_logprobs[0]) < 1:
             distribution_logprobs = None
 

test.py

@@ -5,21 +5,19 @@
 
 #add replicate api key to env
 import os
-os.environ['REPLICATE_API_TOKEN'] = 'r8_aOlrg82Wfg30Rx4L4mv9wI2npPfBQGO0Pvci4'
 
 # def test_decorator(variable_value, prompt_value, context):
 #     return variable_value, prompt_value
 
 async def main():
 
     test = lmql.model(
-        "openai/gpt-3.5-turbo-instruct"
-        # "meta-llama/Llama-2-13b-chat-hf",
-        # endpoint="replicate:deployment/ml-delphai/llama2-13b-chat-lmtp",
+        # "openai/gpt-3.5-turbo-instruct"
+        "meta-llama/Llama-2-13b-chat-hf",
+        endpoint="replicate:deployment/ml-delphai/llama2-13b-chat-lmtp",
         # endpoint="replicate:charles-dyfis-net/llama-2-7b-chat-hf--lmtp-8bit",
-        # tokenizer="AyyYOO/Luna-AI-Llama2-Uncensored-FP16-sharded",
+        tokenizer="AyyYOO/Luna-AI-Llama2-Uncensored-FP16-sharded",
     )
-    pass
 
     answer = await lmql.run(
         """
@@ -32,7 +30,7 @@ def get_probs(variable_value, prompt_value, context):
                 if value > -5:
                     scores[key] = math.exp(value)
             return scores
-        argmax(verbose=True)
+        argmax
         \"How much you like monkeys between 0 and 2?[@get_probs MONKEY]\" where MONKEY in set ([\"0\", \"1\", \"2\"])
         \"How much you like birds between 0 and 2?[@get_probs BIRD]\" where BIRD in set ([\"0\", \"1\", \"2\"])
         return (MONKEY, BIRD)