Views
No views yet
wykonos/moviespip install transformers torch1from transformers import BertTokenizerFast, BertForSequenceClassification
2import torch1model_name = "AventIQ-AI/bert-movie-recommendation-system"
2model = BertForSequenceClassification.from_pretrained(model_name)
3tokenizer = BertTokenizerFast.from_pretrained(model_name)
4
5genre_to_label = {
6 "Action": 0, "Adventure": 1, "Animation": 2, "Comedy": 3, "Crime": 4,
7 "Documentary": 5, "Drama": 6, "Family": 7, "Fantasy": 8, "History": 9,
8 "Horror": 10, "Music": 11, "Mystery": 12, "Romance": 13, "Science Fiction": 14,
9 "TV Movie": 15, "Thriller": 16, "War": 17, "Western": 18
10}
11
12def recommend_movies(genre, top_n=10):
13 """Return a list of movies for a given genre."""
14 if genre not in genre_to_label:
15 return "Unknown Genre"
16 # Filter dataset for movies in the requested genre
17 genre_movies = df[df["genres"].str.contains(genre, case=False, na=False)]["title"].tolist()
18
19 # Return top N movies (or all if fewer exist)
20 return genre_movies[:top_n]
21
22genres_to_test = ["Horror", "Comedy", "Drama"]
23for genre in genres_to_test:
24 recommended_movies = recommend_movies(genre)
25 print(f"Genre: {genre} -> Recommended Movies: {recommended_movies}").
├── model/ # Contains the quantized model files
├── tokenizer_config/ # Tokenizer configuration and vocabulary files
├── model.safetensors/ # Quantized Model
├── README.md # Model documentation