Spaces:
Running
Running
File size: 4,697 Bytes
177b6ad |
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 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 |
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
app = FastAPI()
# Add CORS middleware
app.add_middleware(
CORSMiddleware,
allow_origins=["*"], # In production, replace with your Vercel domain
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Load model and tokenizer
print("Loading SmolLM-135M model...")
MODEL_NAME = "HuggingFaceTB/SmolLM-135M"
# Check for CUDA availability
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")
# Load the model and tokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME).to(device)
print(f"Model loaded successfully! Memory footprint: {model.get_memory_footprint() / 1e6:.2f} MB")
class VibeRequest(BaseModel):
vibe: str
def get_vibe_guidance(vibe: str) -> str:
"""Get specific guidance for each vibe with examples"""
vibe_patterns = {
"romantic": """Generate a romantic and sweet pickup line that's genuine and heartfelt.
Example:
Input: Generate a romantic pickup line
Output: Are you a magician? Because whenever I look at you, everyone else disappears. ❤️
Now generate a romantic pickup line: """,
"cheesy": """Generate a super cheesy and over-the-top pickup line.
Example:
Input: Generate a cheesy pickup line
Output: Are you a parking ticket? Because you've got FINE written all over you! 😏
Now generate a cheesy pickup line: """,
"nerdy": """Generate a nerdy, science-themed pickup line.
Example:
Input: Generate a nerdy pickup line
Output: Are you made of copper and tellurium? Because you're Cu-Te! 🔬
Now generate a nerdy pickup line: """,
"cringe": """Generate the most cringey and over-the-top pickup line imaginable.
Example:
Input: Generate a cringe pickup line
Output: Are you a dictionary? Because you're adding meaning to my life! 📚
Now generate a cringe pickup line: """,
"flirty": """Generate a bold and flirty pickup line.
Example:
Input: Generate a flirty pickup line
Output: Is your name Google? Because you've got everything I've been searching for! 😏
Now generate a flirty pickup line: """
}
return vibe_patterns.get(vibe, "Generate a pickup line with a ")
@app.post("/generate")
async def generate_pickup_line(request: VibeRequest):
try:
vibe = request.vibe
vibe_guide = get_vibe_guidance(vibe)
# Create the prompt
prompt = f"""Instructions: Generate a pickup line with a {vibe} vibe.
{vibe_guide}"""
# Prepare inputs
encoded_input = tokenizer.encode_plus(
prompt,
return_tensors="pt",
padding=True,
return_attention_mask=True
)
input_ids = encoded_input["input_ids"].to(device)
attention_mask = encoded_input["attention_mask"].to(device)
# Generate response
with torch.no_grad():
outputs = model.generate(
input_ids,
attention_mask=attention_mask,
max_new_tokens=100,
do_sample=True,
temperature=0.8,
top_p=0.92,
top_k=50,
pad_token_id=tokenizer.eos_token_id,
eos_token_id=tokenizer.eos_token_id,
)
# Get the full generated text
full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
# Extract just the pickup line
if full_response.startswith(prompt):
response = full_response[len(prompt):].strip()
else:
response = full_response.replace(prompt, "").strip()
# Clean up the response
for marker in ["Instructions:", "Generate a pickup line:", "\n"]:
if marker in response:
response = response.split(marker, 1)[0].strip()
# Add appropriate emoji based on vibe
if vibe == "romantic":
response += " ❤️"
elif vibe == "cheesy":
response += " 😏"
elif vibe == "nerdy":
response += " 🔬"
elif vibe == "cringe":
response += " 😂"
elif vibe == "flirty":
response += " 💋"
return {"pickupLine": response}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8000) |