File size: 830 Bytes
f5b302e |
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 |
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
import json
# Load model
model_name = "google/flan-t5-base"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
rag_pipeline = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
# Load KG
with open("social_graph.json", "r") as f:
kg = json.load(f)
# Build context
person = kg["people"]["bob"]
context = f"Bob is the user's son. They talk about football weekly. Last conversation was about coaching changes."
# User input
query = "What should I say to Bob?"
# RAG-style prompt
prompt = f"""Context: {context}
User wants to say something appropriate to Bob. Suggest a phrase:"""
# Generate
response = rag_pipeline(prompt, max_length=50)
print(response[0]["generated_text"])
|