JohnDoee commited on
Commit
4d93fa1
·
1 Parent(s): 366c5da

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +13 -14
main.py CHANGED
@@ -1,7 +1,7 @@
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
@@ -14,19 +14,20 @@ 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
- mode_name = "johndoee/sentiment" # or "./sentiment"
20
 
21
- # Load the fine-tuned model (johndoee/sentiment) and pass the tokenizer explicitly
 
22
  multilingual_model = pipeline(
23
  "sentiment-analysis",
24
- model= AutoModelForSequenceClassification.from_pretrained(model_name),
25
- tokenizer=original_tokenizer
26
  )
27
 
28
- # Load the English sentiment model
29
- english_model = pipeline("sentiment-analysis", model="siebert/sentiment-roberta-large-english")
30
 
31
  class SentimentRequest(BaseModel):
32
  text: str
@@ -52,11 +53,9 @@ def analyze_sentiment(request: SentimentRequest):
52
  text = request.text
53
  language = detect_language(text)
54
 
55
- # Choose the appropriate model based on language
56
- if language == "en":
57
- result = english_model(text)
58
- else:
59
- result = multilingual_model(text)
60
 
61
  return SentimentResponse(
62
  original_text=text,
 
1
  import os
2
  from fastapi import FastAPI
3
  from pydantic import BaseModel
4
+ from transformers import pipeline, AutoTokenizer
5
  from langdetect import detect, DetectorFactory
6
 
7
  # Ensure consistent language detection results
 
14
 
15
  app = FastAPI()
16
 
17
+ # Model names
18
+ multilingual_model_name = "johndoee/sentiment"
19
+ english_model_name = "siebert/sentiment-roberta-large-english"
20
 
21
+ # Load tokenizer and model for multilingual sentiment analysis
22
+ multilingual_tokenizer = AutoTokenizer.from_pretrained(multilingual_model_name)
23
  multilingual_model = pipeline(
24
  "sentiment-analysis",
25
+ model=multilingual_model_name,
26
+ tokenizer=multilingual_tokenizer
27
  )
28
 
29
+ # Load English sentiment model
30
+ english_model = pipeline("sentiment-analysis", model=english_model_name)
31
 
32
  class SentimentRequest(BaseModel):
33
  text: str
 
53
  text = request.text
54
  language = detect_language(text)
55
 
56
+ # Choose the appropriate model based on detected language
57
+ model = english_model if language == "en" else multilingual_model
58
+ result = model(text)
 
 
59
 
60
  return SentimentResponse(
61
  original_text=text,