JohnDoee commited on
Commit
30d4c48
·
1 Parent(s): be1629b

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +15 -10
main.py CHANGED
@@ -1,28 +1,33 @@
1
  import os
2
- from fastapi import FastAPI
3
  from pydantic import BaseModel
4
- from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
5
  from langdetect import detect, DetectorFactory
 
6
 
7
  # Ensure consistent language detection results
8
  DetectorFactory.seed = 0
9
 
10
- # Set Hugging Face cache directory to a writable location
11
  os.environ["HF_HOME"] = "/tmp/huggingface"
12
  os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface"
13
  os.makedirs(os.environ["HF_HOME"], exist_ok=True)
14
 
15
  app = FastAPI()
16
 
17
- # Load the original tokenizer from the base model
18
  original_tokenizer = AutoTokenizer.from_pretrained("tabularisai/multilingual-sentiment-analysis")
19
 
20
- # Load the fine-tuned model (johndoee/sentiment) and pass the tokenizer explicitly
21
- multilingual_model = pipeline(
22
- "sentiment-analysis",
23
- model="johndoee/sentiment",
24
- tokenizer=original_tokenizer
25
- )
 
 
 
 
26
 
27
  # Load the English sentiment model
28
  english_model = pipeline("sentiment-analysis", model="siebert/sentiment-roberta-large-english")
 
1
  import os
2
+ from fastapi import FastAPI, HTTPException
3
  from pydantic import BaseModel
4
+ from transformers import pipeline, AutoTokenizer
5
  from langdetect import detect, DetectorFactory
6
+ from huggingface_hub import hf_hub_download
7
 
8
  # Ensure consistent language detection results
9
  DetectorFactory.seed = 0
10
 
11
+ # Set Hugging Face cache to a writable location
12
  os.environ["HF_HOME"] = "/tmp/huggingface"
13
  os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface"
14
  os.makedirs(os.environ["HF_HOME"], exist_ok=True)
15
 
16
  app = FastAPI()
17
 
18
+ # Load tokenizer for multilingual model
19
  original_tokenizer = AutoTokenizer.from_pretrained("tabularisai/multilingual-sentiment-analysis")
20
 
21
+ # Dynamically download fine-tuned model (avoid storage issues)
22
+ try:
23
+ MODEL_PATH = hf_hub_download(repo_id="johndoee/sentiment", filename="pytorch_model.bin")
24
+ multilingual_model = pipeline(
25
+ "sentiment-analysis",
26
+ model="johndoee/sentiment",
27
+ tokenizer=original_tokenizer
28
+ )
29
+ except Exception as e:
30
+ raise RuntimeError(f"❌ Error loading model 'johndoee/sentiment': {e}")
31
 
32
  # Load the English sentiment model
33
  english_model = pipeline("sentiment-analysis", model="siebert/sentiment-roberta-large-english")