Spaces:
Running
Running
File size: 1,431 Bytes
9edd84c 81e781d ae8c668 9edd84c c0ea8ea 67165ad e226b59 9a57e7c 00bfe15 9edd84c d4088f1 9edd84c 81e781d 67165ad 81e781d e226b59 |
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 |
from transformers import pipeline
import torch
from langchain.prompts import ChatPromptTemplate
from langchain.schema.output_parser import StrOutputParser
from langchain_huggingface import HuggingFacePipeline
# Initialize HF pipeline for text generation
text_generator = pipeline(
"text-generation", # Task type
model="google/gemma-3n-E2B-it",
# model="google/gemma-3n-e4b-it",
# model="Qwen/Qwen3-Embedding-0.6B",
# device="cuda" if torch.cuda.is_available() else "cpu",
device= "cpu",
torch_dtype=torch.bfloat16,
max_new_tokens=500 # Limit output length
)
# Wrap pipeline for LangChain compatibility
model = HuggingFacePipeline(pipeline=text_generator)
def generate_sentences(topic, n=1):
prompt = ChatPromptTemplate.from_template(
"You are a helpful assistant. Generate exactly {n} simple sentences about the topic: {topic}. "
"Each sentence must be in English and appropriate for all audiences. "
"Return each sentence on a new line without any numbering or bullets"
)
chain = prompt | model | StrOutputParser()
response = chain.invoke({"topic": topic, "n": n})
# Enhanced filtering
return [
line.strip() for line in response.splitlines()
if (line.strip()
and not line.startswith(("###", "Instruction", "Output Format"))
and len(line.split()) <= 15 # Word limit enforcement
)
][:n]
|