Update app.py

app.py

@@ -1,17 +1,12 @@
 import streamlit as st
 from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
-
-# Disable safetensors fast GPU loading (if needed)
-import os
-os.environ["SAFETENSORS_FAST_GPU"] = "0"
 
 # Cache the model and tokenizer
 @st.cache_resource
 def load_model_and_tokenizer():
     model_name = "rajrakeshdr/IntelliSoc"
+    model = AutoModelForCausalLM.from_pretrained(model_name)
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name, use_safetensors=False)
     return model, tokenizer
 
 # Load the model and tokenizer
@@ -29,16 +24,15 @@ if st.button("Generate Text"):
     inputs = tokenizer(prompt, return_tensors="pt", truncation=True, padding=True)
     
     # Generate text
-    with torch.no_grad():
-        outputs = model.generate(
-            inputs.input_ids,
-            max_length=100,
-            num_return_sequences=1,
-            no_repeat_ngram_size=2,
-            top_k=50,
-            top_p=0.95,
-            temperature=0.7
-        )
+    outputs = model.generate(
+        inputs.input_ids,
+        max_length=100,
+        num_return_sequences=1,
+        no_repeat_ngram_size=2,
+        top_k=50,
+        top_p=0.95,
+        temperature=0.7
+    )
     
     # Decode the generated text
     generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)