CoolFace
Apppublic

dgobran/case-study-1

sourceHugging Faceupdated 2y agoView on Hugging Face
0likes
inference.py66 linesDownload Raw Back to root
1from huggingface_hub import InferenceClient2import torch3from transformers import pipeline4 5# Inference client setup6client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")7pipe = pipeline("text-generation", "microsoft/Phi-3-mini-4k-instruct", torch_dtype=torch.bfloat16, device_map="auto")8 9# Global flag to handle cancellation10stop_inference = False11 12# This was adapted from Yang's code and modified to not keep track of history and made compatible with Streamlit13def respond(14    message,15    system_message="You are a bot that paraphrases text.",16    max_tokens=512,17    temperature=0.7,18    top_p=0.95,19    use_local_model=False,20):21    global stop_inference22    stop_inference = False  # Reset cancellation flag23 24    if use_local_model:25        # Local inference26        messages = [{"role": "system", "content": system_message}]27        messages.append({"role": "user", "content": message})28 29        response = ""30        for output in pipe(31            messages,32            max_new_tokens=max_tokens,33            temperature=temperature,34            do_sample=True,35            top_p=top_p,36        ):37            if stop_inference:38                return "Inference cancelled."39            token = output['generated_text'][-1]['content']40            response += token41        42        return response43 44    else:45        # API-based inference46        messages = [{"role": "system", "content": system_message}]47        messages.append({"role": "user", "content": message})48 49        response = ""50        for message_chunk in client.chat_completion(51            messages,52            max_tokens=max_tokens,53            stream=True,54            temperature=temperature,55            top_p=top_p,56        ):57            if stop_inference:58                return "Inference cancelled."59            token = message_chunk.choices[0].delta.content60            response += token61        62        return response63 64def cancel_inference():65    global stop_inference66    stop_inference = True