Done

src/streamlit_app.py

@@ -1,31 +1,21 @@
 import streamlit as st
-import requests
-import os
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
 st.title("Uzmi GPT - Romantic Quote Generator")
 
-# Input prompt
-user_input = st.text_area("Enter your prompt:", "A romantic love quote about forever:")
+@st.cache_resource
+def load_model():
+    tokenizer = AutoTokenizer.from_pretrained("rajan3208/uzmi-gpt")
+    model = AutoModelForCausalLM.from_pretrained("rajan3208/uzmi-gpt")
+    return tokenizer, model
 
-# Button to trigger prediction
-if st.button("Generate"):
-    API_URL = "https://api-inference.huggingface.co/models/rajan3208/uzmi-gpt"
-    headers = {
-        "Authorization": f"Bearer {os.environ.get('HF_TOKEN')}"
-    }
-
-    payload = {"inputs": user_input}
+tokenizer, model = load_model()
 
-    with st.spinner("Generating..."):
-        response = requests.post(API_URL, headers=headers, json=payload)
+prompt = st.text_area("Enter a prompt", "A romantic quote about forever")
 
-    if response.status_code == 200:
-        try:
-            generated_text = response.json()[0]["generated_text"]
-            st.success("Generated Text:")
-            st.write(generated_text)
-        except Exception as e:
-            st.error("Could not parse response.")
-            st.json(response.json())
-    else:
-        st.error(f"Error {response.status_code}: {response.text}")
+if st.button("Generate"):
+    inputs = tokenizer(prompt, return_tensors="pt")
+    output = model.generate(**inputs, max_new_tokens=50)
+    generated = tokenizer.decode(output[0], skip_special_tokens=True)
+    st.success(generated)