File size: 2,313 Bytes
a21df6e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 |
"""
Example script for using NEED AI model from Hugging Face Hub
"""
import torch
from transformers import AutoTokenizer
from huggingface_hub import hf_hub_download
# Configuration
MODEL_REPO = "yogami9/need-ai-conversational-model"
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
def setup_model():
"""Download and setup the model"""
print("📥 Downloading model files...")
# Download model file
model_path = hf_hub_download(
repo_id=MODEL_REPO,
filename="pytorch_model.bin"
)
# Download custom modeling code
modeling_path = hf_hub_download(
repo_id=MODEL_REPO,
filename="modeling_need.py"
)
print("✅ Files downloaded")
# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO)
# Import custom model class
import sys
import os
sys.path.insert(0, os.path.dirname(modeling_path))
from modeling_need import NEEDConversationalModel
# Load model
model = NEEDConversationalModel.from_pretrained(model_path)
model = model.to(DEVICE)
model.eval()
print(f"✅ Model loaded on {DEVICE}")
return model, tokenizer
def generate_response(model, tokenizer, text: str, max_length: int = 100):
"""Generate response for user input"""
# Prepare input
input_text = f"Human: {text}"
input_ids = tokenizer.encode(input_text, return_tensors='pt').to(DEVICE)
speaker_ids = torch.zeros_like(input_ids)
# Generate
with torch.no_grad():
output_ids = model.generate(
input_ids=input_ids,
speaker_ids=speaker_ids,
max_length=max_length,
temperature=0.8,
top_k=50
)
# Decode
response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
return response
if __name__ == "__main__":
# Setup
model, tokenizer = setup_model()
# Example queries
queries = [
"I need a house cleaner in Lagos",
"How much does tutoring cost?",
"I need help with plumbing",
]
print("\n" + "="*60)
print("Testing NEED AI Model")
print("="*60)
for query in queries:
print(f"\n👤 User: {query}")
response = generate_response(model, tokenizer, query)
print(f"🤖 Assistant: {response}")
|