Views
No views yet
transformers and huggingface_hub.pip install huggingface_hub transformers llama-cpp-python1import os
2import glob
3from llama_cpp import Llama
4from transformers import AutoTokenizer
5from huggingface_hub import snapshot_download
6
7
8class VibeThinkerGGUF:
9 def __init__(self, model_source: str = "fdemelo/vibe-thinker-3b-gguf-f16") -> None:
10 """
11 Args:
12 model_source (str): A Hugging Face repo ID (e.g., "fdemelo/vibe-thinker-3b-gguf-f16")
13 OR a path to a local directory (e.g., "./my_model_folder").
14 """
15 if os.path.isdir(model_source):
16 print(f"Found local directory. Loading from: {model_source}")
17 self.model_folder = model_source
18 else:
19 print(f"'{model_source}' is not a local directory. Assuming Hugging Face Repo ID...")
20
21 # Create a sensible local folder name based on the repo ID
22 # e.g., "fdemelo/vibe-thinker-3b-gguf-f16" -> "./vibe-thinker-3b-gguf-f16"
23 repo_name = model_source.split("/")[-1]
24 self.model_folder = f"./{repo_name}"
25
26 print(f"Downloading from repo '{model_source}' to '{self.model_folder}'...")
27 snapshot_download(
28 repo_id=model_source,
29 local_dir=self.model_folder,
30 local_dir_use_symlinks=False
31 )
32
33 gguf_files = glob.glob(os.path.join(self.model_folder, "*.gguf"))
34 if not gguf_files:
35 raise FileNotFoundError(f"No .gguf file was found in {self.model_folder}.")
36
37 model_path = gguf_files[0]
38
39 self.model = Llama(
40 model_path=model_path,
41 n_ctx=4096,
42 n_threads=4,
43 verbose=False
44 )
45
46 self.tokenizer = AutoTokenizer.from_pretrained(
47 self.model_folder,
48 trust_remote_code=True,
49 )
50
51 def infer_text(self, prompt: str) -> str:
52 messages = [{"role": "user", "content": prompt}]
53 text = self.tokenizer.apply_chat_template(
54 messages,
55 tokenize=False,
56 add_generation_prompt=True,
57 )
58
59 response = self.model(
60 prompt=text,
61 max_tokens=4096,
62 temperature=1.0,
63 top_p=0.95,
64 top_k=50,
65 repeat_penalty=1.0,
66 stop=["<|im_end|>", "<|endoftext|>"]
67 )
68
69 return response["choices"][0]["text"]
70
71
72if __name__ == "__main__":
73 model = VibeThinkerGGUF("fdemelo/vibe-thinker-3b-gguf-f16")
74
75 prompt = "Tell me a joke"
76 print(model.infer_text(prompt))
77