Spaces:
Sleeping
Sleeping
File size: 2,434 Bytes
ec77441 1269709 5f743a5 1269709 5f743a5 1269709 ec77441 1269709 ec77441 1269709 ec77441 1269709 ec77441 1269709 ec77441 1269709 ec77441 1269709 ec77441 | 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 | import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import os
from .config import settings
class SyllabusGenerator:
def __init__(self):
self.model = None
self.tokenizer = None
self.is_gpu = settings.DEVICE == "cuda" and torch.cuda.is_available()
self.load_model()
def load_model(self):
model_id = settings.MODEL_ID_GPU if self.is_gpu else settings.MODEL_ID_CPU
print(f"Loading model: {model_id} on {'GPU' if self.is_gpu else 'CPU'}")
self.tokenizer = AutoTokenizer.from_pretrained(model_id, token=settings.HF_TOKEN)
# Load model kwargs
load_kwargs = {
"torch_dtype": torch.bfloat16 if self.is_gpu else torch.float32,
"token": settings.HF_TOKEN
}
if self.is_gpu:
load_kwargs["device_map"] = "auto"
load_kwargs["load_in_4bit"] = True
else:
# For CPU, we rely on standard loading.
# accelerate's device_map="auto" can also be used on CPU but sometimes
# standard loading is more predictable for smaller models if we don't need offloading.
# We'll explicitly move to CPU after load or just let it default.
pass
self.model = AutoModelForCausalLM.from_pretrained(
model_id,
**load_kwargs
)
if not self.is_gpu:
self.model.to("cpu")
def generate(self, prompt: str, context: str = "", max_tokens: int = 1000):
full_prompt = f"""
System: You are an expert academic curriculum designer.
Use the following context to generate a detailed syllabus.
Context: {context}
Question: {prompt}
Syllabus Output:
"""
inputs = self.tokenizer(full_prompt, return_tensors="pt")
if self.is_gpu:
inputs = inputs.to("cuda")
else:
inputs = inputs.to("cpu")
with torch.no_grad():
outputs = self.model.generate(
**inputs,
max_new_tokens=max_tokens,
do_sample=True,
temperature=0.7
)
return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
generator = SyllabusGenerator()
|