update_app.py

app.py

@@ -1,53 +1,44 @@
-# app.py
 import streamlit as st
 import torch
 import openai
 import os
 import time
-import requests
 from PIL import Image
 import tempfile
 import clip  # from OpenAI CLIP repo
 import torch.nn.functional as F
-from transformers import (
-    MBartForConditionalGeneration,
-    MBart50TokenizerFast,
-    AutoTokenizer,
-    AutoModelForCausalLM,
-)
+import requests
+from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
+from transformers import AutoTokenizer, AutoModelForCausalLM
 from rouge_score import rouge_scorer
 from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-openai.api_key = os.getenv("OPENAI_API_KEY")  # Make sure this is set in your environment
+# Set device
+openai.api_key = os.getenv("OPENAI_API_KEY")
 
-# Load MBart model
+# Load MBart
 translator_model = MBartForConditionalGeneration.from_pretrained(
     "facebook/mbart-large-50-many-to-many-mmt",
     device_map="auto",
     low_cpu_mem_usage=True
 )
-translator_tokenizer = MBart50TokenizerFast.from_pretrained(
-    "facebook/mbart-large-50-many-to-many-mmt"
-)
+translator_tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
 translator_tokenizer.src_lang = "ta_IN"
 
-# Load GPT-2 model
-gen_model = AutoModelForCausalLM.from_pretrained(
-    "gpt2",
-    device_map="auto",
-    low_cpu_mem_usage=True
-)
+# Load GPT-2
+gen_model = AutoModelForCausalLM.from_pretrained("gpt2", device_map="auto", low_cpu_mem_usage=True)
 gen_model.eval()
 gen_tokenizer = AutoTokenizer.from_pretrained("gpt2")
 
-# Load CLIP model
+# Load CLIP
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 clip_model, clip_preprocess = clip.load("ViT-B/32", device=device)
 
-# ---- Translation Function ----
+# ---- Translation ----
 def translate_tamil_to_english(text, reference=None):
     start = time.time()
-    inputs = translator_tokenizer(text, return_tensors="pt").to(device)
+    inputs = translator_tokenizer(text, return_tensors="pt")
+    inputs = {k: v.to(translator_model.device) for k, v in inputs.items()}
     outputs = translator_model.generate(
         **inputs,
         forced_bos_token_id=translator_tokenizer.lang_code_to_id["en_XX"]
@@ -63,10 +54,12 @@ def translate_tamil_to_english(text, reference=None):
 
     return translated, duration, rouge_l
 
-# ---- Creative Text Generation ----
+# ---- Creative Text ----
 def generate_creative_text(prompt, max_length=100):
     start = time.time()
-    input_ids = gen_tokenizer.encode(prompt, return_tensors="pt").to(device)
+    input_ids = gen_tokenizer.encode(prompt, return_tensors="pt")
+    input_ids = input_ids.to(gen_model.device)
+
     output = gen_model.generate(
         input_ids,
         max_length=max_length,
@@ -81,14 +74,14 @@ def generate_creative_text(prompt, max_length=100):
     rep_rate = sum(t1 == t2 for t1, t2 in zip(tokens, tokens[1:])) / len(tokens) if len(tokens) > 1 else 0
 
     with torch.no_grad():
-        input_ids = gen_tokenizer.encode(text, return_tensors="pt").to(device)
+        input_ids = gen_tokenizer.encode(text, return_tensors="pt").to(gen_model.device)
         outputs = gen_model(input_ids, labels=input_ids)
         loss = outputs.loss
         perplexity = torch.exp(loss).item()
 
     return text, duration, len(tokens), round(rep_rate, 4), round(perplexity, 4)
 
-# ---- Image Generation using DALL·E 3 ----
+# ---- Image Generation ----
 def generate_image(prompt):
     try:
         start = time.time()
@@ -102,7 +95,6 @@ def generate_image(prompt):
         image_url = response.data[0].url
         image_data = Image.open(requests.get(image_url, stream=True).raw).resize((256, 256))
 
-        # Save to temporary file
         tmp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
         image_data.save(tmp_file.name)
         duration = round(time.time() - start, 2)