fewshot_beam.py

import sys
import requests
import json

context = """[INST] They are planning to host a party next weekend. [/INST] Вони планують провести вечірку наступного вікенду.
[INST] I enjoy swimming in the ocean and feeling the salty breeze. [/INST] Мені подобається плавати в океані та відчувати солоний вітер.
[INST]"""

# docker run --gpus all -p 8000:8000 -e HF_HOME=/hf -e CUDA_VISIBLE_DEVICES=1 -v ~/.cache/huggingface:/hf ghcr.io/mistralai/mistral-src/vllm:latest --host 0.0.0.0 --model mistralai/Mistral-7B-v0.1
def prompt(input, url="http://localhost:8000/v1/completions"):
    data = {
        "prompt": f"{context} {input} [/INST]",
        "stop": "[INST]",
        "max_tokens": 512,
        "temperature": 0,
        #"temperature": 1.0,
        #"top_p": 0.001,
        #"top_k": 40,
        "model": "mistralai/Mistral-7B-v0.1",
        "presence_penalty": 0.1,
        "use_beam_search": True,
        "n": 25,
        "logprobs": 1,
    }
    headers = {
        "Content-Type": "application/json"
    }
    response = requests.post(url, headers=headers, data=json.dumps(data))
    result = response.json()
    return result


for line in sys.stdin:
    text = prompt(line.strip())
    print(json.dumps(text, ensure_ascii=False))