Views
No views yet
1from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
2import torch
3
4# Load model and tokenizer
5model_name = "Abhaykoul/hai3.1-pretrainedv3"
6
7# Set device to CUDA if available
8device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
9
10model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, torch_dtype="auto")
11model.to(device)
12print(model)
13tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
14
15# Message role format for chat
16messages = [
17 {"role": "system", "content": "You are a helpful assistant."},
18 {"role": "user", "content": """hlo"""},
19]
20
21# Apply chat template to format prompt
22prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
23
24# Tokenize input and move to device
25inputs = tokenizer(prompt, return_tensors="pt")
26inputs = {k: v.to(device) for k, v in inputs.items()}
27
28# Set up text streamer for live output
29streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
30
31# Generate text with streaming
32model.generate(
33 **inputs,
34 max_new_tokens=4089,
35 temperature=0.7,
36 top_p=0.9,
37 do_sample=True,
38 streamer=streamer
39)1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer
3
4ckpt = "Abhaykoul/hai3.1-pretrainedv3"
5device = "cuda" if torch.cuda.is_available() else "cpu"
6
7model = AutoModelForCausalLM.from_pretrained(ckpt, trust_remote_code=True).to(device).eval()
8tok = AutoTokenizer.from_pretrained(ckpt, trust_remote_code=True)
9if tok.pad_token is None:
10 tok.pad_token = tok.eos_token
11
12text = "I am thrilled about my new job!"
13enc = tok([text], padding=True, truncation=True, max_length=2048, return_tensors="pt")
14enc = {k: v.to(device) for k, v in enc.items()}
15
16with torch.no_grad():
17 out = model(input_ids=enc["input_ids"], attention_mask=enc.get("attention_mask"), output_hidden_states=True, return_dict=True, use_cache=False)
18 last = out.hidden_states[-1]
19 idx = (enc["attention_mask"].sum(dim=1) - 1).clamp(min=0)
20 pooled = last[torch.arange(last.size(0)), idx]
21 logits = model.structured_lm_head(pooled)
22 pred_id = logits.argmax(dim=-1).item()
23
24print("Predicted class id:", pred_id)
25# Map id -> label using your dataset’s label list, e.g.:
26id2label = ["sadness","joy","love","anger","fear","surprise"] # dair-ai/emotion
27print("Predicted label:", id2label[pred_id] if pred_id < len(id2label) else "unknown")