dgobran/case-study-1
0
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