noddysnots's picture
Update app.py
4df7255 verified
raw
history blame
2.55 kB
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch
import requests
# Load DeepSeek-R1 model with trust_remote_code enabled
model_name = "deepseek-ai/DeepSeek-R1"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
# Ensure compatibility with `flash_attn` and force proper dtype
try:
model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float16, # Forces float16 to prevent fp8 issue
device_map="auto",
trust_remote_code=True
)
except ImportError as e:
raise RuntimeError("Missing required dependency: flash_attn. Install with `pip install flash_attn`") from e
# Use a text-generation pipeline for better inference
generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1)
# 🎯 Function to extract interests from user input
def extract_interests(text):
prompt = f"Extract 3-5 relevant interests from this request: '{text}'. Focus on hobbies and product preferences."
# Generate model output
response = generator(prompt, max_length=50, num_return_sequences=1)
interests = response[0]["generated_text"].replace(prompt, "").strip()
return interests.split(", ") # Convert to a list of keywords
# 🎁 Web search for gift suggestions
def search_gifts(interests):
query = "+".join(interests)
amazon_url = f"https://www.amazon.in/s?k={query}"
flipkart_url = f"https://www.flipkart.com/search?q={query}"
igp_url = f"https://www.igp.com/search?q={query}"
indiamart_url = f"https://dir.indiamart.com/search.mp?ss={query}"
return {
"Amazon": amazon_url,
"Flipkart": flipkart_url,
"IGP": igp_url,
"IndiaMart": indiamart_url
}
# 🎯 Main function for gift recommendation
def recommend_gifts(text):
if not text:
return "Please enter a description."
interests = extract_interests(text) # Extract interests using DeepSeek R1
links = search_gifts(interests) # Get shopping links
return {
"Predicted Interests": interests,
"Gift Suggestions": links
}
# 🎨 Gradio UI for easy interaction
demo = gr.Interface(
fn=recommend_gifts,
inputs="text",
outputs="json",
title="🎁 AI Gift Recommender",
description="Enter details about the person you are buying a gift for, and get personalized suggestions with shopping links!",
)
# πŸš€ Launch Gradio App
if __name__ == "__main__":
demo.launch()