Update app.py

app.py

@@ -1,23 +1,41 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
-# Specify the directory containing the tokenizer's configuration file (config.json)
-model_name = "pytorch_model-00001-of-00002.bin"
+# # Specify the directory containing the tokenizer's configuration file (config.json)
+# model_name = "pytorch_model-00001-of-00002.bin"
 
-# Initialize the tokenizer
-# tokenizer = AutoTokenizer.from_pretrained(model_name, local_files_only=True)
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-tokenizer.pad_token = tokenizer.eos_token
-tokenizer.padding_side = "right"
+# # Initialize the tokenizer
+# # tokenizer = AutoTokenizer.from_pretrained(model_name, local_files_only=True)
+# tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+# tokenizer.pad_token = tokenizer.eos_token
+# tokenizer.padding_side = "right"
 
 
-# Initialize the GPT4All model
+# # Initialize the GPT4All model
+# model = AutoModelForCausalLM.from_pretrained(model_name)
+
+# def generate_text(input_text):
+#     pipe = pipeline(task="text-generation", model=model, tokenizer=tokenizer, max_length=200)
+#     result = pipe(f"<s>[INST] {input_text} [/INST]")
+#     return result[0]['generated_text']
+
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from fastapi import FastAPI
+
+app = FastAPI()
+
+model_name = "pytorch_model-00001-of-00002.bin"  # Replace with your Hugging Face model name
+
 model = AutoModelForCausalLM.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+@app.post("/generate/")
+async def generate_text(prompt: str):
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
+    output = model.generate(input_ids, max_length=50, num_return_sequences=1)
+    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
+    return {"generated_text": generated_text}
 
-def generate_text(input_text):
-    pipe = pipeline(task="text-generation", model=model, tokenizer=tokenizer, max_length=200)
-    result = pipe(f"<s>[INST] {input_text} [/INST]")
-    return result[0]['generated_text']
 
 text_generation_interface = gr.Interface(
     fn=generate_text,